fix the backward for deepspeed (huggingface#9705)

sorami · Jan 20, 2021 · cd5565b · cd5565b
1 parent 538245b
commit cd5565b
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -1282,8 +1282,7 @@ def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor,
             with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                 scaled_loss.backward()
         elif self.deepspeed:
-            # calling on DS engine (model_wrapped == DDP(Deepspeed(PretrainedModule)))
-            self.model_wrapped.module.backward(loss)
+            self.deepspeed.backward(loss)
         else:
             loss.backward()