diff --git a/mergekit/moe/router.py b/mergekit/moe/router.py
index f644e73c..7381739d 100644
--- a/mergekit/moe/router.py
+++ b/mergekit/moe/router.py
@@ -113,15 +113,19 @@ def _do_it(tokenized):
             )
 
     elif mode in ("hidden", "hidden_avg", "hidden_last"):
+        kwargs = {}
+        if load_in_4bit:
+            kwargs["load_in_4bit"] = True
+        if load_in_8bit:
+            kwargs["load_in_8bit"] = True
         model = AutoModelForCausalLM.from_pretrained(
             model_ref.model.path,
             revision=model_ref.model.revision,
             torch_dtype=torch.bfloat16,
             device_map=device,
             low_cpu_mem_usage=True,
-            load_in_4bit=load_in_4bit,
-            load_in_8bit=load_in_8bit,
             trust_remote_code=trust_remote_code,
+            **kwargs,
         )
 
         def _do_it(tokenized):