Follow-up changes to #10575

Author: carmocca

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -180,11 +180,10 @@ def advance(self, data_fetcher: AbstractDataFetcher) -> None:  # type: ignore[ov
 
             # hook
             self.trainer._call_callback_hooks("on_train_batch_start", batch, batch_idx, **extra_kwargs)
-            model_response = self.trainer._call_lightning_module_hook(
+            response = self.trainer._call_lightning_module_hook(
                 "on_train_batch_start", batch, batch_idx, **extra_kwargs
             )
-            ttp_response = self.trainer._call_ttp_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs)
-            response = ttp_response if model_response is None else model_response
+            self.trainer._call_ttp_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs)
             if response == -1:
                 self.batch_progress.increment_processed()
                 raise StopIteration

pytorch_lightning/trainer/trainer.py

@@ -1537,14 +1537,20 @@ def _call_accelerator_hook(
         *args: Any,
         **kwargs: Any,
     ) -> Optional[Any]:
-        self.lightning_module._current_fx_name = hook_name
+        pl_module = self.lightning_module
+        prev_fx_name = pl_module._current_fx_name
+        pl_module._current_fx_name = hook_name
+
         fn = getattr(self.accelerator, hook_name)
         if not callable(fn):
             return None
 
         with self.profiler.profile(hook_name):
             output = fn(*args, **kwargs)
 
+        # restore current_fx when nested context
+        pl_module._current_fx_name = prev_fx_name
+
         return output
 
     @staticmethod