Fix CUDA operations to check device: only use fused ops when input is on CUDA

models/psp/stylegan2/op/fused_act.py

@@ -32,7 +32,7 @@ class FusedLeakyReLUFunctionBackward(Function):
 
         empty = grad_output.new_empty(0)
 
-        if _USE_FUSED:
+        if _USE_FUSED and grad_output.is_cuda:
             grad_input = fused.fused_bias_act(
                 grad_output, empty, out, 3, 1, negative_slope, scale
             )
@@ -65,7 +65,8 @@ class FusedLeakyReLUFunctionBackward(Function):
 class FusedLeakyReLUFunction(Function):
     @staticmethod
     def forward(ctx, input, bias, negative_slope, scale):
-        if _USE_FUSED:
+        # Only use fused operations if CUDA is available and input is on CUDA
+        if _USE_FUSED and input.is_cuda:
             empty = input.new_empty(0)
             out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
         else:

models/psp/stylegan2/op/upfirdn2d.py

@@ -33,7 +33,7 @@ class UpFirDn2dBackward(Function):
 
         grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
 
-        if _USE_EXT:
+        if _USE_EXT and grad_output.is_cuda:
             grad_input = upfirdn2d_op.upfirdn2d(
                 grad_output,
                 grad_kernel,
@@ -76,7 +76,7 @@ class UpFirDn2dBackward(Function):
 
         gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
 
-        if _USE_EXT:
+        if _USE_EXT and gradgrad_input.is_cuda:
             gradgrad_out = upfirdn2d_op.upfirdn2d(
                 gradgrad_input,
                 kernel,
@@ -131,7 +131,7 @@ class UpFirDn2d(Function):
 
         ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
 
-        if _USE_EXT:
+        if _USE_EXT and input.is_cuda:
             out = upfirdn2d_op.upfirdn2d(
                 input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
             )