update: test_sparse_supported

Author: kozistr

tests/test_sparse_gradient.py

@@ -5,6 +5,10 @@
 
 from pytorch_optimizer import load_optimizers
 
+SPARSE_OPTIMIZERS: List[str] = [
+    'madgrad',
+]
+
 NO_SPARSE_OPTIMIZERS: List[str] = [
     'adamp',
     'sgdp',
@@ -32,3 +36,14 @@ def test_sparse_not_supported(no_sparse_optimizer):
 
     with pytest.raises(RuntimeError):
         optimizer.step()
+
+
+@pytest.mark.parametrize('sparse_optimizer', SPARSE_OPTIMIZERS)
+def test_sparse_supported(sparse_optimizer):
+    param = torch.randn(1, 1).to_sparse(1).requires_grad_(True)
+    grad = torch.randn(1, 1).to_sparse(1)
+    param.grad = grad
+
+    optimizer = load_optimizers(optimizer=sparse_optimizer)([param], momentum=0.0)
+    optimizer.zero_grad()
+    optimizer.step()