def test_no_batch_idx_gradient_accumulation():
    """Regression test for an issue where excluding the batch_idx from training_step would disable gradient
    accumulation."""

    class MyModel(BoringModel):
        last_batch_idx = -1

        def training_step(self, batch):  # no batch_idx
            return self.step(batch)

        def optimizer_step(self, epoch, batch_idx, *args, **kwargs):
            assert batch_idx in (1, 3)
            self.last_batch_idx = batch_idx
            return super().optimizer_step(epoch, batch_idx, *args, **kwargs)

    trainer = Trainer(fast_dev_run=4, accumulate_grad_batches=2, limit_val_batches=0)
    model = MyModel()
    trainer.fit(model)
    assert model.last_batch_idx == 3
Regression test for an issue where excluding the batch_idx from training_step would disable gradient
    accumulation.
tests/tests_pytorch/loops/test_training_epoch_loop.py