diff --git a/notebooks/Convolutional layer tests.ipynb b/notebooks/Convolutional layer tests.ipynb
index 8c085cc..f978bfa 100644
--- a/notebooks/Convolutional layer tests.ipynb	
+++ b/notebooks/Convolutional layer tests.ipynb	
@@ -96,10 +96,12 @@
     "        corresponds to `y = conv2d(x, K) + b`.\n",
     "\n",
     "        Args:\n",
-    "            inputs: Array of layer inputs of shape (batch_size, input_dim).\n",
+    "            inputs: Array of layer inputs of shape \n",
+    "                (batch_size, num_input_channels, input_dim_1, input_dim_2).\n",
     "\n",
     "        Returns:\n",
-    "            outputs: Array of layer outputs of shape (batch_size, output_dim).\n",
+    "            outputs: Array of layer outputs of shape \n",
+    "                (batch_size, num_output_channels, output_dim_1, output_dim_2).\n",
     "        \"\"\"\n",
     "        raise NotImplementedError()\n",
     "\n",
@@ -121,7 +123,7 @@
     "\n",
     "        Returns:\n",
     "            Array of gradients with respect to the layer inputs of shape\n",
-    "            (batch_size, input_dim).\n",
+    "            (batch_size, num_input_channels, input_dim_1, input_dim_2).\n",
     "        \"\"\"\n",
     "        raise NotImplementedError()\n",
     "\n",
@@ -132,7 +134,7 @@
     "            inputs: array of inputs to layer of shape (batch_size, input_dim)\n",
     "            grads_wrt_to_outputs: array of gradients with respect to the layer\n",
     "                outputs of shape\n",
-    "                (batch_size, num_output-_channels, output_dim_1, output_dim_2).\n",
+    "                (batch_size, num_output_channels, output_dim_1, output_dim_2).\n",
     "\n",
     "        Returns:\n",
     "            list of arrays of gradients with respect to the layer parameters\n",
@@ -232,8 +234,8 @@
     "           [ 5403.,  5469.,  5535.]]]]\n",
     "    )\n",
     "    layer = layer_class(\n",
-    "        num_input_channels=kernels.shape[0], \n",
-    "        num_output_channels=kernels.shape[1], \n",
+    "        num_input_channels=kernels.shape[1], \n",
+    "        num_output_channels=kernels.shape[0], \n",
     "        input_dim_1=inputs.shape[2], \n",
     "        input_dim_2=inputs.shape[3],\n",
     "        kernel_dim_1=kernels.shape[2],\n",
@@ -319,8 +321,8 @@
     "         [ 226.,  440.,  468.,  222.],\n",
     "         [ 105.,  199.,  209.,   96.]]]])\n",
     "    layer = layer_class(\n",
-    "        num_input_channels=kernels.shape[0], \n",
-    "        num_output_channels=kernels.shape[1], \n",
+    "        num_input_channels=kernels.shape[1], \n",
+    "        num_output_channels=kernels.shape[0], \n",
     "        input_dim_1=inputs.shape[2], \n",
     "        input_dim_2=inputs.shape[3],\n",
     "        kernel_dim_1=kernels.shape[2],\n",
@@ -384,8 +386,8 @@
     "        true_kernel_grads = true_kernel_grads[:, :, ::-1, ::-1]\n",
     "    true_bias_grads = np.array([-126.,   36.])\n",
     "    layer = layer_class(\n",
-    "        num_input_channels=kernels.shape[0], \n",
-    "        num_output_channels=kernels.shape[1], \n",
+    "        num_input_channels=kernels.shape[1], \n",
+    "        num_output_channels=kernels.shape[0], \n",
     "        input_dim_1=inputs.shape[2], \n",
     "        input_dim_2=inputs.shape[3],\n",
     "        kernel_dim_1=kernels.shape[2],\n",