Merge pull request mkocabas#7 from valentinp/gpu_fix

Added casting so that the radius channel works with GPU tensors.
HaoLiMRI · Jul 19, 2018 · d833f1e · d833f1e
2 parents 4279766 + ac4b123
commit d833f1e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/CoordConv.py b/CoordConv.py
@@ -101,7 +101,7 @@ def forward(self, input_tensor):
             yy_channel.type_as(input_tensor)], dim=1)
 
         if self.with_r:
-            rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2))
+            rr = torch.sqrt(torch.pow(xx_channel.type_as(input_tensor) - 0.5, 2) + torch.pow(yy_channel.type_as(input_tensor) - 0.5, 2))
             ret = torch.cat([ret, rr], dim=1)
 
         return ret