update: test cases

Author: kozistr

tests/test_optimizer_parameters.py

@@ -2,18 +2,18 @@
 import torch
 from torch import nn
 
+from pytorch_optimizer.base.exception import NoClosureError
 from pytorch_optimizer.optimizer import (
+    BSAM,
     SAM,
     WSAM,
     Lookahead,
     LookSAM,
-    BSAM,
     PCGrad,
     Ranger21,
     SafeFP16Optimizer,
     load_optimizer,
 )
-from pytorch_optimizer.base.exception import NoClosureError
 from pytorch_optimizer.optimizer.galore_utils import GaLoreProjector
 from tests.constants import PULLBACK_MOMENTUM
 from tests.utils import Example, simple_parameter

tests/test_optimizers.py

@@ -6,15 +6,7 @@
 from torch import nn
 
 from pytorch_optimizer.base.exception import NoClosureError, ZeroParameterSizeError
-from pytorch_optimizer.optimizer import (
-    BSAM,
-    SAM,
-    WSAM,
-    DynamicLossScaler,
-    LookSAM,
-    Muon,
-    load_optimizer,
-)
+from pytorch_optimizer.optimizer import DynamicLossScaler, Muon, load_optimizer
 from pytorch_optimizer.optimizer.alig import l2_projection
 from pytorch_optimizer.optimizer.grokfast import gradfilter_ema, gradfilter_ma
 from pytorch_optimizer.optimizer.scion import build_lmo_norm
@@ -244,12 +236,10 @@ def test_hessian_optimizer(optimizer_name):
     optimizer = load_optimizer(optimizer_name)([param], **parameters)
     optimizer.zero_grad(set_to_none=True)
 
-    # Hutchinson (internal) estimator
     sphere_loss(param).backward(create_graph=True)
     optimizer.step()
     optimizer.zero_grad(set_to_none=True)
 
-    # External estimator
     sphere_loss(param).backward()
     optimizer.step(hessian=torch.zeros_like(param).unsqueeze(0))