Merge pull request #8 from wujiadi1/rc4main

fix the bug of fused_bias_leakyrelu
momo609 · May 29, 2024 · 8c7092b · 8c7092b
2 parents e7c2439 + 18a865b
commit 8c7092b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmcv/ops/fused_bias_leakyrelu.py b/mmcv/ops/fused_bias_leakyrelu.py
@@ -258,7 +258,7 @@ def fused_bias_leakyrelu(input: torch.Tensor,
         torch.Tensor: Feature map after non-linear activation.
     """
 
-    if not input.is_cuda:
+    if not input.is_cuda and input.device.type != 'npu':
         return bias_leakyrelu_ref(input, bias, negative_slope, scale)
 
     return FusedBiasLeakyReLUFunction.apply(input, bias.to(input.dtype),