fixes

Lightning-AI · Sep 15, 2022 · 55ce21b · 55ce21b
1 parent 75e38c4
commit 55ce21b
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/pytorch_lightning/lite/wrappers.py b/src/pytorch_lightning/lite/wrappers.py
@@ -62,9 +62,8 @@ def step(self, closure: Optional[Callable] = None, module: Optional["_LiteModule
         closure = closure or _do_nothing_closure
         return self._strategy.optimizer_step(
             self.optimizer,
-            opt_idx=0,
+            model=(module if module is not None else getattr(self._strategy, "model", None)),
             closure=closure,
-            model=(module if module is not None else getattr(self._strategy, "model", None))
         )