diff --git a/pix2pix-keras.ipynb b/pix2pix-keras.ipynb
index 7ea315c..570a4b6 100644
--- a/pix2pix-keras.ipynb
+++ b/pix2pix-keras.ipynb
@@ -123,9 +123,9 @@
     "       max_layers: max hidden layers\n",
     "    \"\"\"    \n",
     "    if channel_first:\n",
-    "        input_a, input_b =  Input(shape=(nc_in, 256, 256)), Input(shape=(nc_out, 256, 256))\n",
+    "        input_a, input_b =  Input(shape=(nc_in, None, None)), Input(shape=(nc_out, None, None))\n",
     "    else:\n",
-    "        input_a, input_b = Input(shape=(256, 256, nc_in)), Input(shape=(256, 256, nc_out))\n",
+    "        input_a, input_b = Input(shape=(None, None, nc_in)), Input(shape=(None, None, nc_out))\n",
     "    _ = Concatenate(axis=channel_axis)([input_a, input_b])\n",
     "    _ = conv2d(ndf, kernel_size=4, strides=2, padding=\"same\", name = 'First') (_)\n",
     "    _ = LeakyReLU(alpha=0.2)(_)\n",