From 164ff2440d8a4b05db6671f00df22b83bb49436f Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 4 Mar 2024 20:58:21 -0800
Subject: [PATCH] Use the correct PyTorch in the Colab notebook

---
 Colab-TextGen-GPU.ipynb | 5 +++++
 1 file changed, 5 insertions(+)

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",