diff --git a/experimentation/Track model training with MLflow.ipynb b/experimentation/Track model training with MLflow.ipynb
index 6687732..a81458a 100644
--- a/experimentation/Track model training with MLflow.ipynb	
+++ b/experimentation/Track model training with MLflow.ipynb	
@@ -412,7 +412,7 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "Finally, let's try to log an artifact. An artifact can be any file. For example, you can plot the ROC curve and store the plot as an image. The image can be logged as an artifact. \n",
+        "Finally, let's try to log an artifact. An artifact can be any file. For example, you can plot the ROC curve and log the figure as an artifact. \n",
         "\n",
         "Run the following cell to log a parameter, metric, and an artifact."
       ]
@@ -446,11 +446,11 @@
         "    plt.xlabel('False Positive Rate')\n",
         "    plt.ylabel('True Positive Rate')\n",
         "    plt.title('ROC Curve')\n",
-        "    plt.savefig(\"ROC-Curve.png\")\n",
         "\n",
         "    mlflow.log_param(\"estimator\", \"DecisionTreeClassifier\")\n",
-        "    mlflow.log_metric(\"Accuracy\", acc)"
-      ]
+        "    mlflow.log_metric(\"Accuracy\", acc)\n",
+        "    mlflow.log_figure(fig, \"ROC-Curve.png\")"
+      ],
     },
     {
       "attachments": {},