Update test_fused_bias_leakyrelu.py

open-mmlab · Dec 12, 2022 · 7bb7049 · 7bb7049
1 parent 780e3a1
commit 7bb7049
Showing 1 changed file with 7 additions and 27 deletions.
diff --git a/tests/test_ops/test_fused_bias_leakyrelu.py b/tests/test_ops/test_fused_bias_leakyrelu.py
@@ -47,25 +47,12 @@ def test_gradient(self, device):
                     self.input_tensor,
                     delta=1e-4,
                     pt_atol=1e-3)
-            elif IS_NPU_AVAILABLE:
-                gradcheck(
-                    FusedBiasLeakyReLU(2).npu(),
-                    self.input_tensor,
-                    delta=1e-4,
-                    pt_atol=1e-3)
         else:
-            if IS_CUDA_AVAILABLE:
-                gradcheck(
-                    FusedBiasLeakyReLU(2).cuda(),
-                    self.input_tensor,
-                    eps=1e-4,
-                    atol=1e-3)
-            elif IS_NPU_AVAILABLE:
-                gradcheck(
-                    FusedBiasLeakyReLU(2).npu(),
-                    self.input_tensor,
-                    eps=1e-4,
-                    atol=1e-3)
+            gradcheck(
+                FusedBiasLeakyReLU(2).to(device),
+                self.input_tensor,
+                eps=1e-4,
+                atol=1e-3)
 
     @pytest.mark.parametrize('device', [
         pytest.param(
@@ -80,15 +67,8 @@ def test_gradient(self, device):
     def test_gradgradient(self, device):
 
         from mmcv.ops import FusedBiasLeakyReLU
-        if IS_CUDA_AVAILABLE:
-            gradcheck(
-                FusedBiasLeakyReLU(2).cuda(),
-                self.input_tensor,
-                eps=1e-4,
-                atol=1e-3)
-        elif IS_NPU_AVAILABLE:
-            gradcheck(
-                FusedBiasLeakyReLU(2).npu(),
+        gradcheck(
+                FusedBiasLeakyReLU(2).to(device),
                 self.input_tensor,
                 eps=1e-4,
                 atol=1e-3)