diff --git a/coordconv.py b/coordconv.py
index 6d21570..7e21cb3 100644
--- a/coordconv.py
+++ b/coordconv.py
@@ -4,10 +4,11 @@ import torch.nn.modules.conv as conv
 
 
 class AddCoords(nn.Module):
-    def __init__(self, rank, with_r=False):
+    def __init__(self, rank, with_r=False, use_cuda=True):
         super(AddCoords, self).__init__()
         self.rank = rank
         self.with_r = with_r
+        self.use_cuda = use_cuda
 
     def forward(self, input_tensor):
         """
@@ -23,7 +24,7 @@ class AddCoords(nn.Module):
             xx_channel = xx_channel * 2 - 1
             xx_channel = xx_channel.repeat(batch_size_shape, 1, 1)
 
-            if torch.cuda.is_available:
+            if torch.cuda.is_available and self.use_cuda:
                 input_tensor = input_tensor.cuda()
                 xx_channel = xx_channel.cuda()
             out = torch.cat([input_tensor, xx_channel], dim=1)
@@ -57,7 +58,7 @@ class AddCoords(nn.Module):
             xx_channel = xx_channel.repeat(batch_size_shape, 1, 1, 1)
             yy_channel = yy_channel.repeat(batch_size_shape, 1, 1, 1)
 
-            if torch.cuda.is_available:
+            if torch.cuda.is_available and self.use_cuda:
                 input_tensor = input_tensor.cuda()
                 xx_channel = xx_channel.cuda()
                 yy_channel = yy_channel.cuda()
@@ -93,7 +94,7 @@ class AddCoords(nn.Module):
             zx_channel = zx_channel.permute(0, 1, 4, 2, 3)
             zz_channel = torch.cat([zx_channel + i for i in range(dim_y)], dim=3)
 
-            if torch.cuda.is_available:
+            if torch.cuda.is_available and self.use_cuda:
                 input_tensor = input_tensor.cuda()
                 xx_channel = xx_channel.cuda()
                 yy_channel = yy_channel.cuda()