diff --git a/Colab-TextGen-GPU.ipynb b/Colab-TextGen-GPU.ipynb
index 53f60986..82e6c18e 100644
--- a/Colab-TextGen-GPU.ipynb
+++ b/Colab-TextGen-GPU.ipynb
@@ -66,6 +66,11 @@
         "  print(f\"TORCH: {torver}\")\n",
         "  is_cuda118 = '+cu118' in torver  # 2.1.0+cu118\n",
         "\n",
+        "  if is_cuda118:\n",
+        "    !python -m pip install --upgrade torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu118\n",
+        "  else:\n",
+        "    !python -m pip install --upgrade torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121\n",
+        "\n",
         "  textgen_requirements = open('requirements.txt').read().splitlines()\n",
         "  if is_cuda118:\n",
         "      textgen_requirements = [req.replace('+cu121', '+cu118').replace('+cu122', '+cu118') for req in textgen_requirements]\n",