Gagan Bhatia commited on
Commit
4fb639f
1 Parent(s): c8a440e

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +1 -0
src/models/model.py CHANGED
@@ -268,6 +268,7 @@ class Summarization:
268
  elif model_type == "byt5":
269
  self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
270
  self.model = T5ForConditionalGeneration.from_pretrained(
 
271
 
272
  def train(
273
  self,
 
268
  elif model_type == "byt5":
269
  self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
270
  self.model = T5ForConditionalGeneration.from_pretrained(
271
+ f"{model_name}", return_dict=True
272
 
273
  def train(
274
  self,