Fix loss scale (#720)

modelscope · Apr 17, 2024 · 87d24cb · 87d24cb
1 parent 3fecc8c
commit 87d24cb
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/swift/llm/utils/template.py b/swift/llm/utils/template.py
@@ -421,10 +421,10 @@ def data_collator(self,
                                           'constant', 0)
                 labels[0] = F.pad(labels[0], (0, padding_len), 'constant',
                                   -100)
-        if loss_scale:
-            loss_scale[0] = F.pad(loss_scale[0],
-                                  (0, padding_to - labels[0].shape[-1]),
-                                  'constant', 0.)
+                if loss_scale:
+                    loss_scale[0] = F.pad(
+                        loss_scale[0], (0, padding_to - labels[0].shape[-1]),
+                        'constant', 0.)
 
         input_ids = pad_sequence(
             input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
@@ -446,7 +446,7 @@ def data_collator(self,
             'attention_mask': attention_mask,
             'labels': labels,
         }
-        if loss_scale:
+        if loss_scale is not None:
             res['loss_scale'] = loss_scale
         return res