update: test_schedule_free_methods

kozistr · Dec 4, 2024 · 5326483 · 5326483
1 parent e1ab493
commit 5326483
Showing 1 changed file with 6 additions and 12 deletions.
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
@@ -639,18 +639,12 @@ def test_dynamic_scaler():
     scaler.update_scale(overflow=False)
 
 
-def test_schedule_free_train_mode():
-    param = simple_parameter(True)
-
-    opt = load_optimizer('ScheduleFreeAdamW')([param])
-    opt.reset()
-    opt.eval()
-    opt.train()
-
-    opt = load_optimizer('ScheduleFreeSGD')([param])
-    opt.reset()
-    opt.eval()
-    opt.train()
+@pytest.mark.parametrize('optimizer_name', ['ScheduleFreeAdamW', 'ScheduleFreeSGD', 'ScheduleFreeRAdam'])
+def test_schedule_free_methods(optimizer_name):
+    optimizer = load_optimizer(optimizer_name)([simple_parameter(True)])
+    optimizer.reset()
+    optimizer.eval()
+    optimizer.train()
 
 
 @pytest.mark.parametrize('filter_type', ['mean', 'sum'])