Gagan Bhatia commited on
Commit
c59f6db
1 Parent(s): c4b22ed

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +1 -0
src/models/model.py CHANGED
@@ -361,6 +361,7 @@ class Summarization:
361
  )
362
  elif model_type == "mt5":
363
  self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_dir}")
 
364
 
365
  if use_gpu:
366
  if torch.cuda.is_available():
 
361
  )
362
  elif model_type == "mt5":
363
  self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_dir}")
364
+ self.model = MT5ForConditionalGeneration.from_pretrained(
365
 
366
  if use_gpu:
367
  if torch.cuda.is_available():