Merge pull request #328 from kozistr/update/wrapper

[Update] proper property

docs/changelogs/v3.3.5.md

@@ -0,0 +1,9 @@
+### Change Log
+
+### Fix
+
+* Add the missing `state` property in `OrthoGrad` optimizer. (#326, #327)
+
+### Contributions
+
+thanks to @Vectorrent

pytorch_optimizer/optimizer/orthograd.py

@@ -1,17 +1,10 @@
-from collections import defaultdict
 from typing import Callable, Dict
 
 import torch
 from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import (
-    CLOSURE,
-    DEFAULTS,
-    LOSS,
-    OPTIMIZER_INSTANCE_OR_CLASS,
-    STATE,
-)
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER_INSTANCE_OR_CLASS
 
 
 class OrthoGrad(BaseOptimizer):
@@ -27,8 +20,6 @@ def __init__(self, optimizer: OPTIMIZER_INSTANCE_OR_CLASS, **kwargs) -> None:
         self._optimizer_step_post_hooks: Dict[int, Callable] = {}
         self.eps: float = 1e-30
 
-        self.state: STATE = defaultdict(dict)
-
         if isinstance(optimizer, Optimizer):
             self.optimizer = optimizer
         elif 'params' in kwargs:
@@ -46,8 +37,13 @@ def __str__(self) -> str:
     def param_groups(self):
         return self.optimizer.param_groups
 
-    def __getstate__(self):
-        return {'optimizer': self.optimizer}
+    @property
+    def state(self):
+        return self.optimizer.state
+
+    @torch.no_grad()
+    def zero_grad(self) -> None:
+        self.optimizer.zero_grad(set_to_none=True)
 
     @torch.no_grad()
     def reset(self):

requirements-dev.txt

@@ -22,7 +22,7 @@ platformdirs==4.3.6 ; python_version >= "3.8"
 pluggy==1.5.0 ; python_version >= "3.8"
 pytest-cov==5.0.0 ; python_version >= "3.8"
 pytest==8.3.4 ; python_version >= "3.8"
-ruff==0.9.1 ; python_version >= "3.8"
+ruff==0.9.2 ; python_version >= "3.8"
 setuptools==75.8.0 ; python_version >= "3.12"
 sympy==1.12.1 ; python_version == "3.8"
 sympy==1.13.1 ; python_version >= "3.9"

tests/test_gradients.py

@@ -37,7 +37,7 @@ def test_no_gradients(optimizer_name):
     sphere_loss(p1 + p3).backward(create_graph=True)
 
     optimizer.step(lambda: 0.1)  # for AliG optimizer
-    if optimizer_name not in {'lookahead', 'trac'}:
+    if optimizer_name not in {'lookahead', 'trac', 'orthograd'}:
         optimizer.zero_grad(set_to_none=True)
 
 

tests/test_utils.py

@@ -265,6 +265,9 @@ def test_cpu_offload_optimizer():
 
 def test_orthograd_name():
     optimizer = build_orthograd(Example().parameters())
+    optimizer.zero_grad()
+
     _ = optimizer.param_groups
-    _ = optimizer.__getstate__()
+    _ = optimizer.state
+
     assert str(optimizer).lower() == 'orthograd'