Gagan Bhatia commited on
Commit
b4f8dd3
·
1 Parent(s): c22b8b7

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +1 -0
src/models/model.py CHANGED
@@ -355,6 +355,7 @@ class Summarization:
355
  use_gpu (bool, optional): if True, model uses gpu for inferencing/prediction. Defaults to True.
356
  """
357
  if model_type == "t5":
 
358
 
359
  if use_gpu:
360
  if torch.cuda.is_available():
 
355
  use_gpu (bool, optional): if True, model uses gpu for inferencing/prediction. Defaults to True.
356
  """
357
  if model_type == "t5":
358
+ self.tokenizer = T5Tokenizer.from_pretrained(f"{model_dir}")
359
 
360
  if use_gpu:
361
  if torch.cuda.is_available():