bug correction

NVIDIA · Nov 5, 2024 · 5ef5e90 · 5ef5e90
1 parent 8919757
commit 5ef5e90
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/advanced/llm_hf/src/hf_sft_peft_fl.py b/examples/advanced/llm_hf/src/hf_sft_peft_fl.py
@@ -218,7 +218,7 @@ def evaluate(input_weights, mode):
                 out_param["model." + key] = out_param.pop(key).cpu()
 
         # cast out_param to float32 preparing for communication
-        out_param = {k: v.to(torch.float16) for k, v in out_param.items()}
+        out_param = {k: v.to(torch.float32) for k, v in out_param.items()}
 
         # construct trained FL model
         output_model = flare.FLModel(