Avoid deprecation warning after #9901

pytorch_lightning/plugins/training_type/ipu.py

@@ -285,7 +285,7 @@ def on_test_end(self):
     def on_predict_end(self):
         self._detach_models()
 
-    def on_train_batch_start(self, batch: Any, batch_idx: int) -> None:
+    def on_train_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
         # Updates optimizer stats if LR scheduler modified the optimizer state
         optimizer = self.lightning_module.trainer.optimizers[0]
         self.poptorch_models[RunningStage.TRAINING].setOptimizer(optimizer)

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -345,7 +345,7 @@ def on_predict_end(self):
         """Called when predict ends."""
         pass
 
-    def on_train_batch_start(self, batch: Any, batch_idx: int) -> None:
+    def on_train_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
         """Called in the training loop before anything happens for that batch."""
         pass
 

pytorch_lightning/trainer/trainer.py

@@ -1401,14 +1401,14 @@ def call_hook(
             if callable(model_fx):
                 output = model_fx(*args, **kwargs)
 
-            # call the accelerator hook
-            if hook_name not in ("setup", "teardown") and hasattr(self.accelerator, hook_name):
-                accelerator_hook = getattr(self.accelerator, hook_name)
-                accelerator_output = accelerator_hook(*args, **kwargs)
-                # Rely on the accelerator output if lightningModule hook returns nothing
+            # call the ttp hook
+            if hook_name not in ("setup", "teardown") and hasattr(self.training_type_plugin, hook_name):
+                ttp_hook = getattr(self.training_type_plugin, hook_name)
+                ttp_output = ttp_hook(*args, **kwargs)
+                # Rely on the TTP output if lightningModule hook returns nothing
                 # Required for cases such as DataParallel where we reduce the output for the user
                 # todo: move this data parallel logic into the data parallel plugin
-                output = accelerator_output if output is None else output
+                output = ttp_output if output is None else output
 
         if pl_module:
             # restore current_fx when nested context

tests/loops/test_training_loop.py

@@ -86,7 +86,7 @@ def run_training(**trainer_kwargs):
 @pytest.mark.parametrize(["max_epochs", "batch_idx_"], [(2, 5), (3, 8), (4, 12)])
 def test_on_train_batch_start_return_minus_one(max_epochs, batch_idx_, tmpdir):
     class CurrentModel(BoringModel):
-        def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+        def on_train_batch_start(self, batch, batch_idx):
             if batch_idx == batch_idx_:
                 return -1