Update model.py
Browse filesAdd default value to return_dict if not passed.
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,
|