Update pretrain_bert.py (#355)

fix a bug that cause File "pretrain_bert.py", line 91, in loss_func lm_loss_, sop_logits = output_tensor ^^^^^^^^^^^^^^^^^^^^ ValueError: too many values to unpack (expected 2)
microsoft · Feb 27, 2024 · a9856ce · a9856ce
1 parent 31e2584
commit a9856ce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pretrain_bert.py b/pretrain_bert.py
@@ -87,7 +87,7 @@ def data_post_process(data, data_sampler_state_dict):
     return data
 
 def loss_func(loss_mask, sentence_order, output_tensor):
-    lm_loss_, sop_logits = output_tensor
+    lm_loss_, sop_logits, _ = output_tensor
 
     lm_loss_ = lm_loss_.float()
     loss_mask = loss_mask.float()