def test_multiple_optimizers_step(tmp_path):
    """Tests that `step` works with several optimizers."""

    class TestModel(ManualOptModel):
        def training_step(self, batch, batch_idx):
            opt_a, opt_b = self.optimizers()
            x = batch[0]

            loss_1 = self(x)
            loss_1 = self.loss(loss_1, loss_1)

            # make sure there are no grads
            assert_emtpy_grad(self.layer.weight.grad)

            self.manual_backward(loss_1)
            opt_a.step()

            # fake discriminator
            loss_2 = self(x)
            loss_2 = self.loss(loss_2, loss_2)

            # ensure we forward the correct params to the optimizer
            # without retain_graph we can't do multiple backward passes
            self.manual_backward(loss_2, retain_graph=True)
            self.manual_backward(loss_2, retain_graph=True)

            assert self.layer.weight.grad is not None
            opt_b.step()
            opt_b.zero_grad()

            return {"loss1": loss_1.detach(), "loss2": loss_2.detach()}

        # sister test: tests/plugins/test_amp_plugins.py::test_amp_gradient_unscale
        def on_after_backward(self) -> None:
            # check grads are scaled
            scale = self.trainer.precision_plugin.scaler.get_scale()
            assert scale != 1.0  # the return value if not enabled
            grads = [p.grad for p in self.parameters()]
            inv_scale = 1 / scale
            self.original_grads = [p * inv_scale for p in grads]

        def check_grads_unscaled(self, optimizer=None):
            if optimizer is not None:
                scaler = self.trainer.precision_plugin.scaler
                state = scaler._per_optimizer_states[id(optimizer)]
                assert state["stage"].name == "UNSCALED"

            grads = [p.grad for p in self.parameters()]
            assert len(grads) == len(self.original_grads)
            for actual, expected in zip(grads, self.original_grads):
                torch.testing.assert_close(actual, expected)

        def on_before_optimizer_step(self, optimizer, *_):
            self.check_grads_unscaled(optimizer)

    model = TestModel()
    model.val_dataloader = None

    limit_train_batches = 2
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=limit_train_batches,
        limit_val_batches=2,
        max_epochs=1,
        log_every_n_steps=1,
        enable_model_summary=False,
        precision="16-mixed",
        accelerator="gpu",
        devices=1,
    )

    with mock.patch.object(Strategy, "backward", wraps=trainer.strategy.backward) as bwd_mock:
        trainer.fit(model)
    assert bwd_mock.call_count == limit_train_batches * 3
Tests that `step` works with several optimizers.
tests/tests_pytorch/trainer/optimization/test_manual_optimization.py