fix input size with 'with_r' option

HaoLiMRI · Jul 19, 2018 · 7081f3d · 7081f3d
1 parent d833f1e
commit 7081f3d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/CoordConv.py b/CoordConv.py
@@ -112,7 +112,10 @@ class CoordConv(nn.Module):
     def __init__(self, in_channels, out_channels, with_r=False, **kwargs):
         super().__init__()
         self.addcoords = AddCoords(with_r=with_r)
-        self.conv = nn.Conv2d(in_channels + 2, out_channels, **kwargs)
+        in_size = in_channels+2
+        if with_r:
+            in_size += 1
+        self.conv = nn.Conv2d(in_size, out_channels, **kwargs)
 
     def forward(self, x):
         ret = self.addcoords(x)