diff --git a/pyvene_101.ipynb b/pyvene_101.ipynb
index e2b20765..18817bd6 100644
--- a/pyvene_101.ipynb
+++ b/pyvene_101.ipynb
@@ -1877,7 +1877,7 @@
     "    tokenizer(\"The capital of Italy is\", return_tensors=\"pt\"),\n",
     "]\n",
     "base_outputs, counterfactual_outputs = pv_gpt2(\n",
-    "    base, sources, {\"sources->base\": ([[[3]]], [[[3]]])}\n",
+    "    base, sources, {\"sources->base\": ([[[3]]], [[[3]]])}, output_original_output=True\n",
     ")\n",
     "print(counterfactual_outputs.last_hidden_state - base_outputs.last_hidden_state)\n",
     "# call backward will put gradients on model's weights\n",
@@ -2785,7 +2785,7 @@
     "    model=resnet\n",
     ")\n",
     "intervened_outputs = pv_resnet(\n",
-    "    base_inputs, [source_inputs], return_dict=True\n",
+    "    base_inputs, [source_inputs], return_dict=True, output_original_output=True\n",
     ")\n",
     "(intervened_outputs.intervened_outputs.logits - intervened_outputs.original_outputs.logits).sum()"
    ]
@@ -2842,7 +2842,7 @@
     ")\n",
     "\n",
     "intervened_outputs = pv_resnet(\n",
-    "    base_inputs, [source_inputs], return_dict=True\n",
+    "    base_inputs, [source_inputs], return_dict=True, output_original_output=True\n",
     ")\n",
     "(intervened_outputs.intervened_outputs.logits - intervened_outputs.original_outputs.logits).sum()"
    ]