TypeError and other fix

Author: daniellepintz

pytorch_lightning/loops/epoch/evaluation_epoch_loop.py

@@ -218,10 +218,8 @@ def _evaluation_step(self, **kwargs: Any) -> Optional[STEP_OUTPUT]:
             the outputs of the step
         """
         if self.trainer.testing:
-            self.trainer.lightning_module._current_fx_name = "test_step"
             output = self.trainer._call_accelerator_hook("test_step", *kwargs.values())
         else:
-            self.trainer.lightning_module._current_fx_name = "validation_step"
             output = self.trainer._call_accelerator_hook("validation_step", *kwargs.values())
 
         return output

pytorch_lightning/loops/epoch/prediction_epoch_loop.py

@@ -130,7 +130,6 @@ def _predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None
 
         self.batch_progress.increment_started()
 
-        self.trainer.lightning_module._current_fx_name = "predict_step"
         predictions = self.trainer._call_accelerator_hook("predict_step", *step_kwargs.values())
 
         self.batch_progress.increment_processed()

pytorch_lightning/loops/optimization/manual_loop.py

@@ -102,7 +102,6 @@ def advance(self, batch: Any, batch_idx: int) -> None:  # type: ignore[override]
             )
 
             # manually capture logged metrics
-            lightning_module._current_fx_name = "training_step"
             training_step_output = self.trainer._call_accelerator_hook("training_step", *step_kwargs.values())
             self.trainer.training_type_plugin.post_training_step()
 

pytorch_lightning/loops/optimization/optimizer_loop.py

@@ -427,7 +427,6 @@ def _training_step(self, split_batch: Any, batch_idx: int, opt_idx: int) -> Clos
             )
 
             # manually capture logged metrics
-            lightning_module._current_fx_name = "training_step"
             training_step_output = self.trainer._call_accelerator_hook("training_step", *step_kwargs.values())
             self.trainer.training_type_plugin.post_training_step()
 

pytorch_lightning/trainer/trainer.py

@@ -1457,7 +1457,7 @@ def _call_lightning_module_hook(
         pl_module = pl_module or self.lightning_module
 
         if pl_module is None:
-            raise Exception("No Lightning Module is available to call hooks on")
+            raise TypeError("No Lightning Module is available to call hooks on")
 
         fn = getattr(pl_module, hook_name)
         if not callable(fn):
@@ -1538,6 +1538,7 @@ def _call_accelerator_hook(
         *args: Any,
         **kwargs: Any,
     ) -> Optional[Any]:
+        self.lightning_module._current_fx_name = hook_name
         fn = getattr(self.accelerator, hook_name)
         if not callable(fn):
             return None