add unit test

open-mmlab · Mar 20, 2023 · b6834bd · b6834bd
1 parent a384964
commit b6834bd
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/tests/test_ops/test_roi_align.py b/tests/test_ops/test_roi_align.py
@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE
+from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE, IS_NPU_AVAILABLE
 
 _USING_PARROTS = True
 try:
@@ -102,7 +102,11 @@ def _test_roialign_allclose(device, dtype):
     pytest.param(
         'mlu',
         marks=pytest.mark.skipif(
-            not IS_MLU_AVAILABLE, reason='requires MLU support'))
+            not IS_MLU_AVAILABLE, reason='requires MLU support')),
+    pytest.param(
+        'npu',
+        marks=pytest.mark.skipif(
+            not IS_NPU_AVAILABLE, reason='requires NPU support'))
 ])
 @pytest.mark.parametrize('dtype', [
     torch.float,
@@ -111,6 +115,11 @@ def _test_roialign_allclose(device, dtype):
         marks=pytest.mark.skipif(
             IS_MLU_AVAILABLE,
             reason='MLU does not support for 64-bit floating point')),
+    pytest.param(
+        torch.double,
+        marks=pytest.mark.skipif(
+            IS_NPU_AVAILABLE,
+            reason='NPU does not support for 64-bit floating point')),
     torch.half
 ])
 def test_roialign(device, dtype):