update: remove Shampoo from SAM test

Author: kozistr

tests/test_optimizers.py

@@ -167,6 +167,9 @@ def test_sam_optimizers(adaptive, optimizer_sam_config):
     (x_data, y_data), model, loss_fn = build_environment()
 
     optimizer_class, config, iterations = optimizer_sam_config
+    if optimizer_class.__name__ == 'Shampoo':
+        return True
+
     optimizer = SAM(model.parameters(), optimizer_class, **config, adaptive=adaptive)
 
     init_loss, loss = np.inf, np.inf
@@ -190,6 +193,9 @@ def test_sam_optimizers_with_closure(adaptive, optimizer_sam_config):
     (x_data, y_data), model, loss_fn = build_environment()
 
     optimizer_class, config, iterations = optimizer_sam_config
+    if optimizer_class.__name__ == 'Shampoo':
+        return True
+
     optimizer = SAM(model.parameters(), optimizer_class, **config, adaptive=adaptive)
 
     def closure():