Update llama_expo.py (fix bug)

tanganke · Nov 13, 2024 · c7ca936 · c7ca936
1 parent 0d38365
commit c7ca936
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fusion_bench/method/linear/llama_expo.py b/fusion_bench/method/linear/llama_expo.py
@@ -58,7 +58,7 @@ def expo_(
             delta_p = unstructured_magnitude_prune_(
                 delta_p, torch.abs, magnitude_sparsity_ratio, return_pruned_weight=False
             )
-        rlhf_state_dict[n] = rlhf_p + extrapolation_factor * delta_p
+        sft_p = rlhf_p + extrapolation_factor * delta_p
 
         if merge_dtype is not None:
             merged_state_dict[n] = rlhf_p.to(dtype=orignal_dtype)