diff --git a/docs/source/model.ipynb b/docs/source/model.ipynb
index a2b32d8b..0dec9df9 100644
--- a/docs/source/model.ipynb
+++ b/docs/source/model.ipynb
@@ -556,7 +556,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -572,7 +572,7 @@
     "init = torch.rand(batch, NT).log_softmax(-1)\n",
     "\n",
     "dist = torch_struct.SentCFG((terminals, rules, init))\n",
-    "term, rules, init = dist.argmax"
+    "term, rules, init, spans = dist.argmax"
    ]
   },
   {
@@ -1312,9 +1312,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.1"
+   "version": "3.7.4"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }