avichr commited on
Commit
a3b44bc
·
1 Parent(s): 6dd4d02

Update HebEMO.py

Browse files
Files changed (1) hide show
  1. HebEMO.py +2 -2
HebEMO.py CHANGED
@@ -1,5 +1,5 @@
1
  class HebEMO:
2
- def __init__(self, device=0, emotions = ['anticipation', 'joy', 'trust', 'fear', 'surprise', 'anger',
3
  'sadness', 'disgust']):
4
 
5
  from transformers import pipeline
@@ -14,7 +14,7 @@ class HebEMO:
14
  "sentiment-analysis",
15
  model="avichr/hebEMO_"+emo,
16
  tokenizer="avichr/heBERT",
17
- device = self.device #run on GPU
18
  )
19
 
20
 
 
1
  class HebEMO:
2
+ def __init__(self, device=-1, emotions = ['anticipation', 'joy', 'trust', 'fear', 'surprise', 'anger',
3
  'sadness', 'disgust']):
4
 
5
  from transformers import pipeline
 
14
  "sentiment-analysis",
15
  model="avichr/hebEMO_"+emo,
16
  tokenizer="avichr/heBERT",
17
+ device = self.device #-1 run on CPU, else - device ID
18
  )
19
 
20