diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 2dcf1e6a43..c8f4ca1041 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -41,6 +41,13 @@ try:
 except:
     raise ImportError("Pytorch is not installed. Go to https://pytorch.org/.\n"\
                       "We have some installation instructions on our Github page.")
+pass
+
+# Fix up is_bf16_supported https://github.com/unslothai/unsloth/issues/504
+major_version, minor_version = torch.cuda.get_device_capability()
+SUPPORTS_BFLOAT16 = (major_version >= 8)
+def is_bf16_supported(): return SUPPORTS_BFLOAT16
+torch.cuda.is_bf16_supported = is_bf16_supported
 
 # We support Pytorch 2
 # Fixes https://github.com/unslothai/unsloth/issues/38
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index b234a98d8b..c8e00be231 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -16,6 +16,7 @@ from dataclasses import dataclass, field
 from typing import Optional
 from transformers import TrainingArguments
 from trl import SFTTrainer
+from . import is_bfloat16_supported
 
 __all__ = [
     "UnslothTrainingArguments",