diff --git a/01.train_ODC_XGBoost.ipynb b/01.train_ODC_XGBoost.ipynb
index 691d333..e938a06 100644
--- a/01.train_ODC_XGBoost.ipynb
+++ b/01.train_ODC_XGBoost.ipynb
@@ -3000,10 +3000,13 @@
     "from sklearn.metrics import ConfusionMatrixDisplay\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
+    "# Create figure first\n",
+    "fig, ax = plt.subplots(figsize=(10, 8))\n",
+    "\n",
     "class_names = list(label_mapping.keys())\n",
     "cm = confusion_matrix(y_test_np, y_pred_test)\n",
     "disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names)\n",
-    "disp.plot(figsize=(10, 8), cmap='Blues')\n",
+    "disp.plot(cmap='Blues', ax=ax)\n",
     "plt.xticks(rotation=45, ha='right')\n",
     "plt.title('XGBoost Confusion Matrix')\n",
     "plt.tight_layout()\n",