Lolalb commited on
Commit
5424c8e
·
verified ·
1 Parent(s): 2e41a1b

Update model.py

Browse files

Add default value to return_dict if not passed.

Files changed (1) hide show
  1. model.py +2 -1
model.py CHANGED
@@ -378,7 +378,8 @@ class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
378
  labels: Optional[torch.Tensor] = None,
379
  return_dict: Optional[bool] = None,
380
  ):
381
-
 
382
  output = self.model.forward(
383
  input_ids=input_ids,
384
  position_ids=position_ids,
 
378
  labels: Optional[torch.Tensor] = None,
379
  return_dict: Optional[bool] = None,
380
  ):
381
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
382
+
383
  output = self.model.forward(
384
  input_ids=input_ids,
385
  position_ids=position_ids,