Fix notebook error (#15)

Author: wohlert

Mental rotation.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 150,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -77,9 +77,13 @@
     "f, axarr = plt.subplots(1, 15, figsize=(20, 7))\n",
     "\n",
     "x, v = next(iter(loader))\n",
+    "x = x.squeeze(0)\n",
+    "\n",
+    "scene_id = 0\n",
     "\n",
     "for i, ax in enumerate(axarr.flat):\n",
-    "    ax.imshow(x[0][i].permute(1, 2, 0))\n",
+    "    # Move channel dimension to end\n",
+    "    ax.imshow(x[scene_id][i].permute(1, 2, 0))\n",
     "    ax.axis(\"off\")"
    ]
   },
@@ -302,13 +306,6 @@
     "for i, ax in enumerate(axarr.flat):\n",
     "    ax.imshow(x_samples[i].data.permute(1, 2, 0))"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -327,7 +324,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.2"
+   "version": "3.6.0"
   }
  },
  "nbformat": 4,