From fdfe1f59f56935f1945269e5beda50969810158a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 31 Jul 2024 08:54:22 -0700
Subject: [PATCH] Update _utils.py

---
 unsloth/models/_utils.py | 40 ++++++++++++++++++++--------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index c9bc6065..fe3aa904 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -192,20 +192,20 @@ def patch_mistral_nemo_config(config):
 # Get Xformers
 from xformers import __version__ as xformers_version
 # Temporarily disable 0.0.27 and higher - inference issues
-# if Version(xformers_version) >= Version("0.0.27"):
-#     raise ImportError(
-#         "Unsloth: If you are in Colab, we updated the top cell install instructions - please change it to below "\
-#         "then press Disconnect Runtime and then Restart it.\n"\
-#         "\n"\
-#         "%%capture\n"
-#         "# Installs Unsloth, Xformers (Flash Attention) and all other packages!\n"
-#         '!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n'
-#         '!pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes\n'\
-#         '\n'\
-#         f"Otherwise in local machines, your xformers version of {xformers_version} is too new.\n"\
-#         'Please downgrade xformers via `pip install --force-reinstall "xformers<0.0.27"'
-#     )
-# pass
+if Version(xformers_version) >= Version("0.0.27"):
+    raise ImportError(
+        "Unsloth: If you are in Colab, we updated the top cell install instructions - please change it to below "\
+        "then press Disconnect Runtime and then Restart it.\n"\
+        "\n"\
+        "%%capture\n"
+        "# Installs Unsloth, Xformers (Flash Attention) and all other packages!\n"
+        '!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n'
+        '!pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes\n'\
+        '\n'\
+        f"Otherwise in local machines, your xformers version of {xformers_version} is too new.\n"\
+        'Please downgrade xformers via `pip install --force-reinstall "xformers<0.0.27"'
+    )
+pass
 
 if   Version(torch_version) < Version("2.2.0") and Version(xformers_version) >= Version("0.0.24"):
     raise ImportError(
@@ -217,12 +217,12 @@ def patch_mistral_nemo_config(config):
         f"Unsloth: You have torch = {torch_version} but xformers = {xformers_version}.\n"\
         f"Please install xformers < 0.0.26 for torch = {torch_version}."
     )
-# elif Version(torch_version) < Version("2.4.0") and Version(xformers_version) >= Version("0.0.27"):
-#     raise ImportError(
-#         f"Unsloth: You have torch = {torch_version} but xformers = {xformers_version}.\n"\
-#         f"Please install xformers < 0.0.27 for torch = {torch_version}."
-#     )
-# pass
+elif Version(torch_version) < Version("2.4.0") and Version(xformers_version) >= Version("0.0.27"):
+    raise ImportError(
+        f"Unsloth: You have torch = {torch_version} but xformers = {xformers_version}.\n"\
+        f"Please install xformers < 0.0.27 for torch = {torch_version}."
+    )
+pass
 
 from xformers._cpp_lib import _register_extensions
 try: