diff --git a/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb b/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb
index b988b83..efd0183 100644
--- a/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb
+++ b/PyTorch/Profiling_and_Optimization/Profiler_and_Optimization.ipynb
@@ -140,6 +140,17 @@
     "The collected trace files will be saved to ./hpu_profile"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c97a9142-a80e-4c4f-82b7-4f08509996c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from optimum.habana.utils import HabanaProfile\n",
+    "HabanaProfile.enable()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 5,
@@ -202,7 +213,7 @@
     }
    ],
    "source": [
-    "!python run_image_classification.py \\\n",
+    "%run run_image_classification.py \\\n",
     "    --model_name_or_path microsoft/swin-base-patch4-window7-224-in22k \\\n",
     "    --dataset_name cifar10 \\\n",
     "    --output_dir /tmp/outputs/ \\\n",
@@ -233,7 +244,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
+   "id": "4f029263-187b-46be-b017-09c5be1ef8dc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mv ./hpu_profile/*.pt.trace.json ./hpu_profile/UNOPT.pt.trace.json"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
    "id": "f23106fe-e911-4e67-9c8f-4ad8d587408e",
    "metadata": {},
    "outputs": [
@@ -251,8 +272,7 @@
     }
    ],
    "source": [
-    "%cd hpu_profile\n",
-    "%ls -al"
+    "%ls -al hpu_profile"
    ]
   },
   {
@@ -303,7 +323,7 @@
    ],
    "source": [
     "%load_ext tensorboard\n",
-    "%tensorboard --logdir=~/optimum-habana/examples/image-classification/hpu_profile --port 6006    # Your port selection may vary, default is 6006"
+    "%tensorboard --logdir=./hpu_profile --port 6006    # Your port selection may vary, default is 6006"
    ]
   },
   {
@@ -348,7 +368,7 @@
     }
    ],
    "source": [
-    "!habana_perf_tool --trace /root/optimum-habana/examples/image-classification/hpu_profile/sc09wynn05-hls2_14734.1729284340533778439.pt.trace.json"
+    "!habana_perf_tool --trace ./hpu_profile/UNOPT.pt.trace.json"
    ]
   },
   {
@@ -381,9 +401,18 @@
     }
    ],
    "source": [
-    "%cd ..\n",
-    "!hl-prof-config -e off -phase=multi-enq -g 1-20 -s my_profiling_session\n",
-    "!export HABANA_PROFILE=1"
+    "!hl-prof-config -e off -phase=multi-enq -g 1-20 -s my_profiling_session"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1eb9fe7f-5817-4803-90b8-244da14213c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"HABANA_PROFILE\"] = \"1\""
    ]
   },
   {
@@ -393,7 +422,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!HABANA_PROFILE=1 python run_image_classification.py \\\n",
+    "%run run_image_classification.py \\\n",
     "    --model_name_or_path microsoft/swin-base-patch4-window7-224-in22k \\\n",
     "    --dataset_name cifar10 \\\n",
     "    --output_dir /tmp/outputs/ \\\n",