diff --git a/training/scripts/train.py b/training/scripts/train.py
index b99aaee..bde830a 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -89,7 +89,8 @@ def train(config_path):
         save_total_limit=config["train_params"]["save_total_limit"],
         eval_strategy=config.get("eval_strategy", "steps"),
         eval_steps=config.get("eval_steps", 100),
-        mixed_precision=config.get("mixed_precision", "bf16"),
+        bf16=config.get("mixed_precision", "bf16") == "bf16",
+        fp16=config.get("mixed_precision", "bf16") == "fp16",
         gradient_checkpointing=config.get("gradient_checkpointing", True),
     )