diff --git a/Train-Runpod.ipynb b/Train-Runpod.ipynb
index 48371e1..cad4060 100644
--- a/Train-Runpod.ipynb
+++ b/Train-Runpod.ipynb
@@ -73,6 +73,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "!pip install --upgrade jupyterlab\n",
+    "!pip install --upgrade ipywidgets\n",
     "!pip install -q omegaconf\n",
     "!pip install -q einops\n",
     "!pip install -q pytorch-lightning==1.6.5\n",
@@ -86,7 +88,6 @@
     "!pip install -q torchmetrics==0.6.0\n",
     "!pip install -e .\n",
     "!pip install huggingface_hub\n",
-    "!pip install ipywidgets==7.7.1\n",
     "import ipywidgets as widgets\n",
     "print(finish_msg)"
    ]