From a8ddc39482a39f13a4dce2458ee15370fd7f797c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 6 Nov 2024 19:00:13 -0800
Subject: [PATCH] Update loader.py

---
 unsloth/models/loader.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 4566302ed0..0414d54460 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -54,11 +54,11 @@ def _get_dtype(dtype):
         "bfloat16": torch.bfloat16,
         torch.bfloat16: torch.bfloat16,
     }
-    if dtype in __DTYPE_MAP:
-        return __DTYPE_MAP[dtype]
+    if dtype is None or dtype == None: return None
+    if dtype in __DTYPE_MAP: return __DTYPE_MAP[dtype]
     else:
-        print(f"Unsloth: {dtype} is not recognized, so we'll default to torch.float16")
-        return torch.float16
+        print(f"Unsloth: {dtype} is not recognized, so we'll default to None")
+        return None
     pass
 pass