diff --git a/mu_map/models/discriminator.py b/mu_map/models/discriminator.py
index dd90d3562404e628822ec160108fec236f9209da..e1fe36586ec946b0b00cd8d034854ce90dae16e4 100644
--- a/mu_map/models/discriminator.py
+++ b/mu_map/models/discriminator.py
@@ -85,7 +85,7 @@ class Discriminator(nn.Module):
         self.conv = nn.Sequential(*conv_layers)
 
         # input is halved by each convolutional layer
-        self.fc_input_size = map(lambda x: x // 2 ** (len(conv_features)), input_size)
+        self.fc_input_size = map(lambda x: x // 2 ** (len(conv_features)), self.input_size)
         self.fc_input_size = reduce(lambda x, y: x * y, self.fc_input_size)
 
         fc_layers = []
@@ -249,6 +249,7 @@ if __name__ == "__main__":
         help="batch size of inputs for a test computation",
     )
     args = parser.parse_args()
+    args.input_size = tuple(args.input_size)
 
     if args.type == "class":
         net = Discriminator(