fix CausalLM
Browse files- modeling_ltgbert.py +1 -1
modeling_ltgbert.py
CHANGED
@@ -528,7 +528,7 @@ class LtgbertForCausalLM(LtgbertModel):
|
|
528 |
)
|
529 |
return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
|
530 |
|
531 |
-
return
|
532 |
loss=masked_lm_loss,
|
533 |
logits=subword_prediction,
|
534 |
hidden_states=contextualized_embeddings if output_hidden_states else None,
|
|
|
528 |
)
|
529 |
return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
|
530 |
|
531 |
+
return CausalLMOutput(
|
532 |
loss=masked_lm_loss,
|
533 |
logits=subword_prediction,
|
534 |
hidden_states=contextualized_embeddings if output_hidden_states else None,
|