Update init_deepspeed api

microsoft · Apr 2, 2021 · 1506ee5 · 1506ee5
1 parent 3a66e31
commit 1506ee5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -911,7 +911,7 @@ def train(
         if self.args.deepspeed:
             if self.args.ort:
                 self.model = model
-            model, optimizer, lr_scheduler = init_deepspeed(self, model, num_training_steps=max_steps)
+            model, optimizer, lr_scheduler = init_deepspeed(self, num_training_steps=max_steps)
             self.model = model.module._original_module if self.args.ort else model.module
             self.model_wrapped = model  # will get further wrapped in DDP
             self.deepspeed = model  # DeepSpeedEngine object