diff --git a/coursework3.ipynb b/coursework3.ipynb
index 829601a..b1aa536 100644
--- a/coursework3.ipynb
+++ b/coursework3.ipynb
@@ -158,7 +158,7 @@
     "        .groupby(['doc_class', 'word_id']) \\\n",
     "        ['count'].apply(sum) \\\n",
     "        .unstack(fill_value=0)\n",
-    "    occurrences = np.zeros((K, len(V)))\n",
+    "    occurrences = np.zeros((K, W))\n",
     "    occurrences[x.index.values.reshape((-1,1)), x.columns.values] = x\n",
     "    word_count = np.sum(occurrences, axis=1)\n",
     "    \n",