From 7081f3dd740dcdb049dec7b84d27065fa8ddeacb Mon Sep 17 00:00:00 2001
From: Muhammed Kocabas <muhammed.kocabas@metu.edu.tr>
Date: Thu, 19 Jul 2018 23:07:59 +0300
Subject: [PATCH] fix input size with 'with_r' option

---
 CoordConv.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/CoordConv.py b/CoordConv.py
index 33da5df..d2e7cdd 100644
--- 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)