update: test_sam_family_methods

Author: kozistr

tests/test_optimizer_parameters.py

@@ -7,11 +7,13 @@
     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
@@ -86,16 +88,21 @@ def test_lookahead_parameters():
         Lookahead(optimizer, pullback_momentum='invalid')
 
 
-@pytest.mark.parametrize('optimizer', [SAM, WSAM, LookSAM])
+@pytest.mark.parametrize('optimizer', [SAM, WSAM, LookSAM, BSAM])
 def test_sam_family_methods(optimizer):
     base_optimizer = load_optimizer('lion')
 
-    opt = optimizer(params=[simple_parameter()], model=None, base_optimizer=base_optimizer)
-    opt.init_group({})
+    opt = optimizer(params=[simple_parameter()], model=None, base_optimizer=base_optimizer, num_data=1)
+    opt.zero_grad()
+
+    opt.init_group({'params': []})
     opt.load_state_dict(opt.state_dict())
 
+    with pytest.raises(NoClosureError):
+        opt.step()
+
     with pytest.raises(ValueError):
-        optimizer(model=None, params=None, base_optimizer=base_optimizer, rho=-0.1)
+        optimizer(model=None, params=None, base_optimizer=base_optimizer, rho=-0.1, num_data=1)
 
 
 def test_safe_fp16_methods():

tests/test_optimizers.py

@@ -210,34 +210,6 @@ def test_closure(optimizer):
         optimizer.step(closure=dummy_closure)
 
 
-def test_no_closure():
-    param = simple_parameter()
-
-    optimizer = SAM([param], load_optimizer('adamp'))
-    optimizer.zero_grad()
-
-    with pytest.raises(NoClosureError):
-        optimizer.step()
-
-    optimizer = WSAM(None, [param], load_optimizer('adamp'))
-    optimizer.zero_grad()
-
-    with pytest.raises(NoClosureError):
-        optimizer.step()
-
-    optimizer = BSAM([param], 1)
-    optimizer.zero_grad()
-
-    with pytest.raises(NoClosureError):
-        optimizer.step()
-
-    optimizer = LookSAM([param], load_optimizer('adamp'))
-    optimizer.zero_grad()
-
-    with pytest.raises(NoClosureError):
-        optimizer.step()
-
-
 def test_nero_zero_scale():
     param = simple_parameter()