simplify casting to device and dtype

axolotl-ai-cloud · Dec 29, 2023 · 097ce6c · 097ce6c
1 parent 7767d7d
commit 097ce6c
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py
@@ -103,9 +103,7 @@ def do_inference(
             importlib.import_module("axolotl.prompters"), prompter
         )
 
-    model = model.to(cfg.device)
-    if cfg.bf16:
-        model = model.to(torch.bfloat16)
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
     while True:
         print("=" * 80)
@@ -170,9 +168,7 @@ def do_inference_gradio(
             importlib.import_module("axolotl.prompters"), prompter
         )
 
-    model = model.to(cfg.device)
-    if cfg.bf16:
-        model = model.to(torch.bfloat16)
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
     def generate(instruction):
         if not instruction: