fix quick-start notebook

Author: rromb

Diff for: scripts/taming-transformers.ipynb

@@ -122,7 +122,7 @@
     "outputId": "e6e1ce0f-bce9-4ebe-a852-95003ca7b630"
    },
    "source": [
-    "%pip install omegaconf>=2.0.0 pytorch-lightning>=1.0.8\n",
+    "%pip install omegaconf>=2.0.0 pytorch-lightning>=1.0.8 einops transformers\n",
     "import sys\n",
     "sys.path.append(\".\")"
    ],
@@ -461,7 +461,7 @@
     "c_code, c_indices = model.encode_to_c(segmentation)\n",
     "print(\"c_code\", c_code.shape, c_code.dtype)\n",
     "print(\"c_indices\", c_indices.shape, c_indices.dtype)\n",
-    "assert c_code.shape[2]*c_code.shape[3] == c_indices.shape[1]\n",
+    "assert c_code.shape[2]*c_code.shape[3] == c_indices.shape[0]\n",
     "segmentation_rec = model.cond_stage_model.decode(c_code)\n",
     "show_segmentation(torch.softmax(segmentation_rec, dim=1))"
    ],