def _optimizer_to_device(optimizer: Optimizer, device: _DEVICE) -> None:
    """Moves the state of a single optimizer to the device."""
    for p, v in optimizer.state.items():
        if not isinstance(v, MutableMapping):
            # Support for custom optimizers
            optimizer.state[p] = apply_to_collection(v, Tensor, move_data_to_device, device, allow_frozen=True)
            continue
        for key, val in v.items():
            # The 'step' parameter needs to remain unmoved (possibly on the CPU) since that is where the optimizer
            # needs it. See https://github.com/pytorch/pytorch/issues/74424
            if key != "step":
                v[key] = move_data_to_device(val, device)
Moves the state of a single optimizer to the device.
src/lightning/fabric/utilities/optimizer.py