cointegrated commited on
Commit
928013e
1 Parent(s): f44876d

change the loading args

Browse files
Files changed (1) hide show
  1. translation.py +1 -1
translation.py CHANGED
@@ -100,7 +100,7 @@ def sentenize_with_fillers(text, splitter, fix_double_space=True, ignore_errors=
100
 
101
  class Translator:
102
  def __init__(self):
103
- self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_URL, low_cpu_mem_usage=False, load_in_8bit=True)
104
  if torch.cuda.is_available():
105
  self.model.cuda()
106
  self.tokenizer = NllbTokenizer.from_pretrained(MODEL_URL)
 
100
 
101
  class Translator:
102
  def __init__(self):
103
+ self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_URL, low_cpu_mem_usage=True)
104
  if torch.cuda.is_available():
105
  self.model.cuda()
106
  self.tokenizer = NllbTokenizer.from_pretrained(MODEL_URL)