def test_gradient_accumulation_scheduling_last_batch(tmp_path, accumulate_grad_batches, limit_train_batches):
    """Verify optimizer.step() applied to last batch while grad accumulation."""

    class TestModel(BoringModel):
        def state_dict(self, *args, **kwargs):
            return deepcopy(super().state_dict(*args, **kwargs))

        def check(self, d1, d2, equal=True):
            keys = d1.keys() | d2.keys()
            values = [torch.equal(d1[k], d2[k]) for k in keys]
            return all(values) if equal else not any(values)

        def backward(self, *args, **kwargs) -> None:
            pre_bwd_state_dict = self.state_dict()
            assert self.check(self.start_state_dict, pre_bwd_state_dict)

            out = super().backward(*args, **kwargs)

            # state dict is equal, just the gradients changed
            assert self.check(pre_bwd_state_dict, self.state_dict())

            return out

        def optimizer_step(self, *args, **kwargs):
            pre_opt_step_state_dict = self.state_dict()
            assert self.check(self.start_state_dict, pre_opt_step_state_dict)

            # this calls `backward` and `on_after_backward` inside the closure
            out = super().optimizer_step(*args, **kwargs)

            # the state dict changed
            assert self.check(pre_opt_step_state_dict, self.state_dict(), equal=False)

            self.opt_step_called = True
            return out

        def on_train_batch_start(self, *_):
            self.start_state_dict = self.state_dict()
            self.opt_step_called = False

        def on_train_batch_end(self, outputs, batch, batch_idx):
            end_state_dict = self.state_dict()
            is_last_batch = (batch_idx + 1) == self.trainer.num_training_batches

            if is_last_batch or self.opt_step_called:
                assert self.check(self.start_state_dict, end_state_dict, equal=False)
            else:
                assert self.check(self.start_state_dict, end_state_dict)

    model = TestModel()
    trainer = Trainer(
        accumulate_grad_batches=accumulate_grad_batches,
        max_epochs=2,
        limit_train_batches=limit_train_batches,
        limit_val_batches=0,
        default_root_dir=tmp_path,
        enable_progress_bar=False,
    )

    trainer.fit(model)
Verify optimizer.step() applied to last batch while grad accumulation.
tests/tests_pytorch/trainer/test_trainer.py