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

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +1 -0
src/models/model.py CHANGED
@@ -356,6 +356,7 @@ class Summarization:
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():
 
356
  """
357
  if model_type == "t5":
358
  self.tokenizer = T5Tokenizer.from_pretrained(f"{model_dir}")
359
+ self.model = T5ForConditionalGeneration.from_pretrained(
360
 
361
  if use_gpu:
362
  if torch.cuda.is_available():