diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 6144efe485..af7e1eb293 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -151,10 +151,10 @@ if hasattr(transformers.cache_utils, "DynamicCache") and \
     source = source.split("\n")
     source = "\n".join(x[start:] for x in source)
     where = source.find("raise KeyError")
-    # source = source[:where] + \
-    #     f"if len(self) == 0:\n{spaces}{spaces}"\
-    #     "    raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \
-    #     f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:]
+    source = source[:where] + \
+        f"if len(self) == 0:\n{spaces}{spaces}"\
+        "    raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \
+        f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:]
     source = source.replace("__getitem__", "__cache_utils_getitem__", 1)
     exec(source)
     transformers.cache_utils.DynamicCache.__getitem__ = __cache_utils_getitem__