Fix: bf16 support for inference (#981)

* Fix: bf16 torch dtype * simplify casting to device and dtype --------- Co-authored-by: Wing Lian <wing.lian@gmail.com>
axolotl-ai-cloud · Dec 29, 2023 · 3678a6c · 3678a6c
1 parent f8ae59b
commit 3678a6c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py
@@ -103,7 +103,7 @@ def do_inference(
             importlib.import_module("axolotl.prompters"), prompter
         )
 
-    model = model.to(cfg.device)
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
     while True:
         print("=" * 80)
@@ -168,7 +168,7 @@ def do_inference_gradio(
             importlib.import_module("axolotl.prompters"), prompter
         )
 
-    model = model.to(cfg.device)
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
     def generate(instruction):
         if not instruction: