fix OrthoGrad state management bug (#327)

* fix state management bug

* run black with proper args

* execute ruff

Author: Vectorrent

pytorch_optimizer/optimizer/orthograd.py

@@ -1,10 +1,17 @@
+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
+from pytorch_optimizer.base.types import (
+    CLOSURE,
+    DEFAULTS,
+    LOSS,
+    OPTIMIZER_INSTANCE_OR_CLASS,
+    STATE,
+)
 
 
 class OrthoGrad(BaseOptimizer):
@@ -20,6 +27,8 @@ 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: