AttributeError: 'OrthoGrad' object has no attribute 'state'

[Vectorrent]

Describe the bug

I am trying to test the new OrthoGrad wrapper, but it fails with the following error, when I try to use it with AdamW and Lion (and presumably others). It appears that Pytorch Lightning expects an optimizer.state attribute to exist on the OrthoGrad optimizer; however, it does not, and thus crashes.

I tested Lookahead, and it's working fine. OrthoGrad is the only wrapper giving me trouble.

To Reproduce

• OS : Arch Linux
• PyTorch version : 2.5.1
• PyTorch-Lightning : 2.5.0.post0
• Python version : 3.13.1
• pytorch-optimizer version : 3.3.4

Log

Traceback (most recent call last):
  File "/home/crow/repos/praxis/run.py", line 1113, in <module>
    trainer.fit(
    ~~~~~~~~~~~^
        train_model,
        ^^^^^^^^^^^^
        datamodule,
        ^^^^^^^^^^^
        ckpt_path=ckpt_path,
        ^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/crow/repos/praxis/.venv/lib/python3.13/site-packages/lightning/pytorch/trainer/trainer.py", line 539, in fit
    call._call_and_handle_interrupt(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/crow/repos/praxis/.venv/lib/python3.13/site-packages/lightning/pytorch/trainer/call.py", line 68, in _call_and_handle_interrupt
    trainer._teardown()
    ~~~~~~~~~~~~~~~~~^^
  File "/home/crow/repos/praxis/.venv/lib/python3.13/site-packages/lightning/pytorch/trainer/trainer.py", line 1005, in _teardown
    self.strategy.teardown()
    ~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/crow/repos/praxis/.venv/lib/python3.13/site-packages/lightning/pytorch/strategies/strategy.py", line 532, in teardown
    _optimizers_to_device(self.optimizers, torch.device("cpu"))
    ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/crow/repos/praxis/.venv/lib/python3.13/site-packages/lightning/fabric/utilities/optimizer.py", line 27, in _optimizers_to_device
    _optimizer_to_device(opt, device)
    ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
  File "/home/crow/repos/praxis/.venv/lib/python3.13/site-packages/lightning/fabric/utilities/optimizer.py", line 32, in _optimizer_to_device
    for p, v in optimizer.state.items():
                ^^^^^^^^^^^^^^^
AttributeError: 'OrthoGrad' object has no attribute 'state'

Expected behavior

I would expect the wrapper to behave like Lookahead does: as a wrapper around an Optimizer instance.

[Vectorrent]

This PR should address the bug. I've tested locally, and it works.
#327