ensure merged model matches the training dtype (#902)

* ensure merged model matches the training dtype * Update src/axolotl/cli/__init__.py * Update src/axolotl/cli/__init__.py
axolotl-ai-cloud · Nov 29, 2023 · 1d21aa6 · 1d21aa6
1 parent 71b7ea3
commit 1d21aa6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/cli/__init__.py b/src/axolotl/cli/__init__.py
@@ -72,7 +72,7 @@ def do_merge_lora(
 
     LOG.info("running merge of LoRA with base model")
     model = model.merge_and_unload()
-    model.to(dtype=torch.float16)
+    model.to(dtype=cfg.torch_dtype)
 
     if cfg.local_rank == 0:
         LOG.info(f"saving merged model to: {str(Path(cfg.output_dir) / 'merged')}")