test: update tests

XuehaiPan · XuehaiPan · commit 7737f70c15f7 · 2023-02-05T14:03:15.000+08:00
diff --git a/tests/test_accelerated_op.py b/tests/test_accelerated_op.py
@@ -21,6 +21,47 @@
 import torchopt
 
 
+try:
+    import torchopt._C.adam_op
+except ImportError:
+    CXX_ACCELERATED_OP_AVAILABLE = False
+else:
+    CXX_ACCELERATED_OP_AVAILABLE = True
+
+
+def test_accelerated_op_is_available():
+    assert torchopt.accelerated_op_available('cpu')
+    assert torchopt.accelerated_op_available(torch.device('cpu'))
+
+    if CXX_ACCELERATED_OP_AVAILABLE:
+        assert not torchopt.accelerated_op_available('meta')
+        assert not torchopt.accelerated_op_available(torch.device('meta'))
+        assert not torchopt.accelerated_op_available(['cpu', 'meta'])
+        assert not torchopt.accelerated_op_available([torch.device('cpu'), torch.device('meta')])
+    else:
+        assert torchopt.accelerated_op_available('meta')
+        assert torchopt.accelerated_op_available(torch.device('meta'))
+        assert torchopt.accelerated_op_available(['cpu', 'meta'])
+        assert torchopt.accelerated_op_available([torch.device('cpu'), torch.device('meta')])
+
+    if torch.cuda.is_available():
+        assert torchopt.accelerated_op_available()
+        assert torchopt.accelerated_op_available('cuda')
+        assert torchopt.accelerated_op_available('cuda:0')
+        assert torchopt.accelerated_op_available(0)
+        assert torchopt.accelerated_op_available(['cpu', 'cuda'])
+        assert torchopt.accelerated_op_available(['cpu', 'cuda:0'])
+        assert torchopt.accelerated_op_available(['cpu', 0])
+    else:
+        assert not torchopt.accelerated_op_available()
+        assert not torchopt.accelerated_op_available('cuda')
+        assert not torchopt.accelerated_op_available('cuda:0')
+        assert not torchopt.accelerated_op_available(0)
+        assert not torchopt.accelerated_op_available(['cpu', 'cuda'])
+        assert not torchopt.accelerated_op_available(['cpu', 'cuda:0'])
+        assert not torchopt.accelerated_op_available(['cpu', 0])
+
+
 @helpers.parametrize(
     dtype=[torch.float64, torch.float32],
     lr=[1e-2, 1e-3, 1e-4],
diff --git a/torchopt/accelerated_op/__init__.py b/torchopt/accelerated_op/__init__.py
@@ -41,5 +41,5 @@ def is_available(devices: Optional[Union[Device, Iterable[Device]]] = None) -> b
             updates = torch.tensor(1.0, device=device)
             op(updates, updates, updates, 1)
         return True
-    except BaseException:  # pylint: disable=broad-except
+    except Exception:  # pylint: disable=broad-except
         return False
