fix compute dtype

Former-commit-id: 8e3540c62d748dcecbcb3a29d57abe665e7b31c8
2025-08-22 22:02:51 +08:00 · 2023-07-05 15:13:00 +08:00 · 2023-07-05 15:13:00 +08:00 · 7ba52f5b6e
commit 7ba52f5b6e
parent c87910ada3
1 changed files with 1 additions and 1 deletions
--- a/src/utils/common.py
+++ b/src/utils/common.py
@ -213,7 +213,7 @@ def load_pretrained(
    model = AutoModelForCausalLM.from_pretrained(
        model_to_load,
        config=config,
-        torch_dtype=model_args.compute_dtype,
+        torch_dtype=torch.bfloat16 if model_args.compute_dtype == torch.bfloat16 else torch.float16,
        low_cpu_mem_usage=True,
        **config_kwargs
    )