wassname/CoordConv
1
0
Fork 0
mirror of https://github.com/wassname/CoordConv.git synced 2026-09-10 11:40:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added manual control over cuda usage

Browse Source
This commit is contained in:
bakirillov
2019-10-18 13:58:37 +03:00
parent 3f8626458f
commit 482e550797
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coordconv.py
+5 -4
View File
@@ -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()