diff --git a/train.ipynb b/train.ipynb
index 9cc2f4b..a7becc3 100644
--- a/train.ipynb
+++ b/train.ipynb
@@ -117,7 +117,7 @@
     "    \n",
     "    bert = AlbertModel(albert_base_configuration)\n",
     "    bert = MultiTaskModel(bert, \n",
-    "                          num_vocab=max([m['token'] for m in token_maps.values()]), \n",
+    "                          num_vocab=1 + max([m['token'] for m in token_maps.values()]), \n",
     "                          num_tokens=config['model_params']['vocab_size'],\n",
     "                          hidden_size=config['model_params']['hidden_size'])\n",
     "    \n",