Apply suggestions from code review

Author: Borda

src/lightning/pytorch/core/module.py

@@ -1161,8 +1161,9 @@ def training_step(...):
                     opt.step()
 
         """
+        self.toggle_optimizer(optimizer)
         try:
-            yield self.toggle_optimizer(optimizer)
+            yield
         finally:
             self.untoggle_optimizer(optimizer)