def test_skip_training_step_with_grad_scaler():
    """Test that the grad scaler gets skipped when skipping a training step."""

    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            if batch_idx % 2:
                return None  # skipping the backward should skip the grad scaler too
            return super().training_step(batch, batch_idx)

    trainer = Trainer(
        accelerator="cuda",
        devices=1,
        precision="16-mixed",
        barebones=True,
        max_steps=5,
        gradient_clip_val=0.5,
    )
    assert isinstance(trainer.precision_plugin, MixedPrecision)
    assert trainer.precision_plugin.scaler is not None
    trainer.precision_plugin.scaler = Mock(wraps=trainer.precision_plugin.scaler)
    model = TestModel()
    trainer.fit(model)
    assert trainer.precision_plugin.scaler.unscale_.call_count == 3
    assert trainer.precision_plugin.scaler.step.call_count == 3
    assert trainer.precision_plugin.scaler.update.call_count == 3
Test that the grad scaler gets skipped when skipping a training step.
tests/tests_pytorch/plugins/precision/test_amp_integration.py