diff --git a/tutorials/train_model.ipynb b/tutorials/train_model.ipynb
index a41fbb1e..a74101e5 100644
--- a/tutorials/train_model.ipynb
+++ b/tutorials/train_model.ipynb
@@ -251,18 +251,15 @@
    ],
    "source": [
     "# Install required packages.\n",
-    "import os\n",
-    "import torch\n",
-    "\n",
-    "!pip install pyg-lib -f https://data.pyg.org/whl/torch-2.3.0+cpu.html # PyG for working with graphs\n",
-    "!pip install git+https://github.com/pyg-team/pytorch_geometric.git # more PyG\n",
-    "!pip install pytorch_frame[full] #PyTorch Frame for working with tabular data\n",
-    "!pip install relbench[full]"
+    "!pip install torch==2.4.0\n",
+    "!pip install torch-geometric torch-sparse torch-scatter torch-cluster torch-spline-conv pyg-lib -f https://data.pyg.org/whl/torch-2.4.0+cpu.html\n",
+    "!pip install pytorch_frame\n",
+    "!pip install relbench"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2024-07-26T00:05:04.397024Z",
@@ -285,6 +282,8 @@
     }
    ],
    "source": [
+    "import os\n",
+    "import torch\n",
     "import relbench\n",
     "\n",
     "relbench.__version__"
@@ -565,8 +564,8 @@
      "text": [
       "Collecting sentence-transformers\n",
       "  Downloading sentence_transformers-3.0.1-py3-none-any.whl (227 kB)\n",
-      "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/227.1 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r",
-      "\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m \u001b[32m225.3/227.1 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r",
+      "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/227.1 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\n",
+      "\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m \u001b[32m225.3/227.1 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\n",
       "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m227.1/227.1 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
       "\u001b[?25hRequirement already satisfied: transformers<5.0.0,>=4.34.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (4.41.2)\n",
       "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (4.66.4)\n",