More fix

arcee-ai · Dec 30, 2023 · 2bd639e · 2bd639e
1 parent 7e9af2d
commit 2bd639e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mergekit/scripts/mixtral_moe.py b/mergekit/scripts/mixtral_moe.py
@@ -188,7 +188,9 @@ def build(
     if config.dtype:
         out_dtype = dtype_from_name(config.dtype)
     elif base_cfg.torch_dtype:
-        out_dtype = dtype_from_name(base_cfg.torch_dtype)
+        out_dtype = base_cfg.torch_dtype
+        if isinstance(out_dtype, str):
+            out_dtype = dtype_from_name(out_dtype)
     else:
         out_dtype = None