refactor: ZeroParameterSizeError

Author: kozistr

pytorch_optimizer/optimizer/adai.py

@@ -4,7 +4,7 @@
 from torch.optim.optimizer import Optimizer
 
 from pytorch_optimizer.base.base_optimizer import BaseOptimizer
-from pytorch_optimizer.base.exception import NoSparseGradientError, ZeroParameterSize
+from pytorch_optimizer.base.exception import NoSparseGradientError, ZeroParameterSizeError
 from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 
@@ -124,7 +124,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 exp_avg_sq_hat_sum += exp_avg_sq.sum() / bias_correction2
 
         if param_size == 0:
-            raise ZeroParameterSize()
+            raise ZeroParameterSizeError()
 
         exp_avg_sq_hat_mean = exp_avg_sq_hat_sum / param_size
 

pytorch_optimizer/optimizer/ranger21.py

@@ -6,7 +6,7 @@
 from torch.optim import Optimizer
 
 from pytorch_optimizer.base.base_optimizer import BaseOptimizer
-from pytorch_optimizer.base.exception import NoSparseGradientError, ZeroParameterSize
+from pytorch_optimizer.base.exception import NoSparseGradientError, ZeroParameterSizeError
 from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.gc import centralize_gradient
@@ -227,7 +227,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         # stable weight decay
         if param_size == 0:
-            raise ZeroParameterSize()
+            raise ZeroParameterSizeError()
 
         variance_normalized = math.sqrt(variance_ma_sum / param_size)
         if math.isnan(variance_normalized):

tests/test_optimizer_parameters.py

@@ -5,7 +5,7 @@
 from torch import nn
 
 from pytorch_optimizer import SAM, Lookahead, PCGrad, Ranger21, SafeFP16Optimizer, load_optimizer
-from pytorch_optimizer.base.exception import ZeroParameterSize
+from pytorch_optimizer.base.exception import ZeroParameterSizeError
 from tests.utils import Example
 
 OPTIMIZER_NAMES: List[str] = [
@@ -219,7 +219,7 @@ def test_size_of_parameter(optimizer):
     model: nn.Module = nn.Linear(1, 1, bias=False)
     model.requires_grad_(False)
 
-    with pytest.raises(ZeroParameterSize):
+    with pytest.raises(ZeroParameterSizeError):
         load_optimizer(optimizer)(model.parameters(), 100).step()
 
 

tests/test_optimizers.py

@@ -30,7 +30,7 @@
     SafeFP16Optimizer,
     Shampoo,
 )
-from pytorch_optimizer.base.exception import ZeroParameterSize
+from pytorch_optimizer.base.exception import ZeroParameterSizeError
 from tests.utils import (
     MultiHeadLogisticRegression,
     build_environment,
@@ -361,7 +361,7 @@ def test_closure(optimizer):
 
     try:
         optimizer.step(closure=dummy_closure)
-    except ZeroParameterSize:  # in case of Ranger21, Adai optimizers
+    except ZeroParameterSizeError:  # in case of Ranger21, Adai optimizers
         pass