def test_hooks_with_different_argument_names(tmp_path):
    """Test that argument names can be anything in the hooks."""

    class CustomBoringModel(BoringModel):
        def assert_args(self, x, batch_nb):
            assert isinstance(x, Tensor)
            assert x.size() == (1, 32)
            assert isinstance(batch_nb, int)

        def training_step(self, x1, batch_nb1):
            self.assert_args(x1, batch_nb1)
            return super().training_step(x1, batch_nb1)

        def validation_step(self, x2, batch_nb2):
            self.assert_args(x2, batch_nb2)
            return super().validation_step(x2, batch_nb2)

        # we don't support a different name for `dataloader_idx`
        def test_step(self, x3, batch_nb3, dataloader_idx):
            self.assert_args(x3, batch_nb3)
            assert isinstance(dataloader_idx, int)
            return super().test_step(x3, batch_nb3)

        # we don't support a different name for `dataloader_idx`
        def predict_step(self, x4, batch_nb4, dataloader_idx):
            self.assert_args(x4, batch_nb4)
            assert isinstance(dataloader_idx, int)
            return super().predict_step(x4, batch_nb4, dataloader_idx)

        def test_dataloader(self):
            return [DataLoader(RandomDataset(32, 64)), DataLoader(RandomDataset(32, 64))]

        def predict_dataloader(self):
            return [DataLoader(RandomDataset(32, 64)), DataLoader(RandomDataset(32, 64))]

    model = CustomBoringModel()

    trainer = Trainer(devices=1, default_root_dir=tmp_path, fast_dev_run=5)

    trainer.fit(model)
    trainer.test(model)

    preds = trainer.predict(model)
    assert len(preds) == 2
    assert all(len(x) == 5 for x in preds)
Test that argument names can be anything in the hooks.
tests/tests_pytorch/models/test_hooks.py