Update xbert.py

salesforce · Mar 1, 2022 · b726f8e · b726f8e
1 parent f224b67
commit b726f8e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/xbert.py b/models/xbert.py
@@ -1297,7 +1297,7 @@ def forward(
             lm_loss = lm_loss.view(prediction_scores.size(0),-1).sum(1)
 
         if soft_labels is not None:
-            loss_distill = -torch.sum(F.log_softmax(shifted_prediction_scores, dim=1)*soft_labels,dim=-1)
+            loss_distill = -torch.sum(F.log_softmax(shifted_prediction_scores, dim=-1)*soft_labels,dim=-1)
             loss_distill = (loss_distill * (labels!=-100)).sum(1)
             lm_loss = (1-alpha)*lm_loss + alpha*loss_distill                    
 
@@ -1426,7 +1426,7 @@ def forward(
             masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
 
         if soft_labels is not None:
-            loss_distill = -torch.sum(F.log_softmax(prediction_scores, dim=1)*soft_labels,dim=-1)
+            loss_distill = -torch.sum(F.log_softmax(prediction_scores, dim=-1)*soft_labels,dim=-1)
             loss_distill = loss_distill[labels!=-100].mean()
             masked_lm_loss = (1-alpha)*masked_lm_loss + alpha*loss_distill