diff --git a/module-4/sub-graph.ipynb b/module-4/sub-graph.ipynb
index 54ab17ac..9c81b06b 100644
--- a/module-4/sub-graph.ipynb
+++ b/module-4/sub-graph.ipynb
@@ -166,7 +166,7 @@
     "    fa_summary = \"Poor quality retrieval of Chroma documentation.\"\n",
     "    return {\"fa_summary\": fa_summary, \"processed_logs\": [f\"failure-analysis-on-log-{failure['id']}\" for failure in failures]}\n",
     "\n",
-    "fa_builder = StateGraph(input=FailureAnalysisState,output=FailureAnalysisOutputState)\n",
+    "fa_builder = StateGraph(FailureAnalysisState,output=FailureAnalysisOutputState)\n",
     "fa_builder.add_node(\"get_failures\", get_failures)\n",
     "fa_builder.add_node(\"generate_summary\", generate_summary)\n",
     "fa_builder.add_edge(START, \"get_failures\")\n",
@@ -226,7 +226,7 @@
     "    report = \"foo bar baz\"\n",
     "    return {\"report\": report}\n",
     "\n",
-    "qs_builder = StateGraph(input=QuestionSummarizationState,output=QuestionSummarizationOutputState)\n",
+    "qs_builder = StateGraph(QuestionSummarizationState,output=QuestionSummarizationOutputState)\n",
     "qs_builder.add_node(\"generate_summary\", generate_summary)\n",
     "qs_builder.add_node(\"send_to_slack\", send_to_slack)\n",
     "qs_builder.add_edge(START, \"generate_summary\")\n",