Gagan Bhatia commited on
Commit
26d8504
1 Parent(s): cf4737f

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +1 -0
src/models/model.py CHANGED
@@ -258,6 +258,7 @@ class Summarization:
258
  if model_type == "t5":
259
  self.tokenizer = T5Tokenizer.from_pretrained(f"{model_name}")
260
  self.model = T5ForConditionalGeneration.from_pretrained(
 
261
 
262
  def train(
263
  self,
 
258
  if model_type == "t5":
259
  self.tokenizer = T5Tokenizer.from_pretrained(f"{model_name}")
260
  self.model = T5ForConditionalGeneration.from_pretrained(
261
+ f"{model_name}", return_dict=True
262
 
263
  def train(
264
  self,