Scale loss before backward (#35207)

huggingface · Dec 23, 2024 · 3cd3cd5 · 3cd3cd5
1 parent f5264a8
commit 3cd3cd5
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -3698,10 +3698,12 @@ def training_step(
             with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                 scaled_loss.backward()
         else:
-            self.accelerator.backward(loss, **kwargs)
             # Finally we need to normalize the loss for reporting
             if num_items_in_batch is None:
-                return loss.detach() / self.args.gradient_accumulation_steps
+                loss /= self.args.gradient_accumulation_steps
+
+            self.accelerator.backward(loss, **kwargs)
+
             return loss.detach()
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):