update: test cases

Author: kozistr

tests/test_gradients.py

@@ -2,7 +2,7 @@
 import torch
 
 from pytorch_optimizer.base.exception import NoSparseGradientError
-from pytorch_optimizer.optimizer import SAM, TRAC, WSAM, AdamP, Lookahead, OrthoGrad, load_optimizer
+from pytorch_optimizer.optimizer import SAM, TRAC, WSAM, AdamP, Lookahead, LookSAM, OrthoGrad, load_optimizer
 from tests.constants import NO_SPARSE_OPTIMIZERS, SPARSE_OPTIMIZERS, VALID_OPTIMIZER_NAMES
 from tests.utils import build_environment, simple_parameter, simple_sparse_parameter, sphere_loss
 
@@ -116,12 +116,13 @@ def test_sparse_supported(sparse_optimizer):
             optimizer.step()
 
 
-def test_sam_no_gradient():
+@pytest.mark.parametrize('optimizer', [SAM, LookSAM])
+def test_sam_no_gradient(optimizer):
     (x_data, y_data), model, loss_fn = build_environment()
     model.fc1.weight.requires_grad = False
     model.fc1.weight.grad = None
 
-    optimizer = SAM(model.parameters(), AdamP)
+    optimizer = optimizer(model.parameters(), AdamP)
     optimizer.zero_grad()
 
     loss = loss_fn(y_data, model(x_data))

tests/test_optimizers.py

@@ -232,6 +232,30 @@ def test_looksam_optimizer(environment):
     assert tensor_to_numpy(init_loss) > 2.0 * tensor_to_numpy(loss)
 
 
+def test_looksam_optimizer_with_closure(environment):
+    (x_data, y_data), model, loss_fn = environment
+
+    optimizer = LookSAM(model.parameters(), load_optimizer('adamw'), lr=5e-1)
+
+    def closure():
+        first_loss = loss_fn(y_data, model(x_data))
+        first_loss.backward()
+        return first_loss
+
+    init_loss, loss = np.inf, np.inf
+    for _ in range(5):
+        loss = loss_fn(y_data, model(x_data))
+        loss.backward()
+
+        optimizer.step(closure)
+        optimizer.zero_grad()
+
+        if init_loss == np.inf:
+            init_loss = loss
+
+    assert tensor_to_numpy(init_loss) > 2.0 * tensor_to_numpy(loss)
+
+
 @pytest.mark.parametrize('adaptive', ADAPTIVE_FLAGS)
 @pytest.mark.parametrize('decouple', DECOUPLE_FLAGS)
 def test_wsam_optimizer(adaptive, decouple, environment):