Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/trainer/optimization/test_manual_optimization.py (Python)

Code before:
        limit_val_batches=0,
        precision="16-mixed",
        accelerator="gpu",
        devices=1,
    )
    trainer.fit(model)


class CustomMapping(collections.abc.Mapping):
    """A custom implementation of Mapping for testing purposes."""

    def __init__(self, *args, **kwargs):
        self._store = dict(*args, **kwargs)

    def __getitem__(self, key):
        return self._store[key]

    def __iter__(self):
        return iter(self._store)

    def __len__(self):
        return len(self._store)

    def __repr__(self):
        return f"{self.__class__.__name__}({self._store})"


@RunIf(min_cuda_gpus=1)
@pytest.mark.parametrize("dicttype", [dict, CustomMapping])
def test_multiple_optimizers_step(tmp_path, dicttype):
    """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

Review comment (suggestion):
```suggestion
    def __repr__(self):
        return f"{self.__class__.__name__}({self._store})"

    def __copy__(self):
        cls = self.__class__
        new_obj = cls(self._store.copy())
        return new_obj
```

Seems like we also need a copy method here to make sure copy works in [`.../manual.py`](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/loops/optimization/manual.py#L49)

Code after:
        limit_val_batches=0,
        precision="16-mixed",
        accelerator="gpu",
        devices=1,
    )
    trainer.fit(model)


class CustomMapping(collections.abc.Mapping):
    """A custom implementation of Mapping for testing purposes."""

    def __init__(self, *args, **kwargs):
        self._store = dict(*args, **kwargs)

    def __getitem__(self, key):
        return self._store[key]

    def __iter__(self):
        return iter(self._store)

    def __len__(self):
        return len(self._store)

    def __repr__(self):
        return f"{self.__class__.__name__}({self._store})"

    def __copy__(self):
        cls = self.__class__
        new_obj = cls(self._store.copy())
        return new_obj

    def copy(self):
        return self.__copy__()


@RunIf(min_cuda_gpus=1)
@pytest.mark.parametrize("dicttype", [dict, CustomMapping])
def test_multiple_optimizers_step(tmp_path, dicttype):
    """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)