davda54 commited on
Commit
ea5705f
1 Parent(s): 3d07f17

fix CausalLM

Browse files
Files changed (1) hide show
  1. 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 MaskedLMOutput(
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,