Update src/axolotl/utils/models.py

Co-authored-by: Wing Lian <wing.lian@gmail.com>
axolotl-ai-cloud · Dec 22, 2023 · 68982f8 · 68982f8
1 parent a546d81
commit 68982f8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -307,7 +307,7 @@ def load_model(
         device_map = infer_auto_device_map(
             model_canvas,
             max_memory=max_memory,
-            dtype="float16",  # TODO: may probably use bfloat16 and others here as well
+            dtype=cfg.torch_dtype,
         )
         # We can discard max_memory now as we have a device map set up for us
         max_memory = None