Fix: bf16 torch dtype

axolotl-ai-cloud · Dec 19, 2023 · 959845e · 959845e
1 parent d25c34c
commit 959845e
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py
@@ -112,6 +112,8 @@ def do_inference(
         )
 
     model = model.to(cfg.device)
+    if cfg.bf16:
+        model = model.to(torch.bfloat16)
 
     while True:
         print("=" * 80)
@@ -185,6 +187,8 @@ def do_inference_gradio(
         )
 
     model = model.to(cfg.device)
+    if cfg.bf16:
+        model = model.to(torch.bfloat16)
 
     def generate(instruction):
         if not instruction: