Fix processing client weights

SymbioticLab · Dec 17, 2023 · c3a940a · c3a940a
1 parent 26a1f8c
commit c3a940a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/fedscale/cloud/aggregation/optimizers.py b/fedscale/cloud/aggregation/optimizers.py
@@ -98,8 +98,10 @@ def update_round_gradient(
                 ) + (1.0 / learning_rate) * np.float_power(loss + 1e-10, qfedq)
 
             # update global model
-            for idx, param in enumerate(target_model.parameters()):
-                param.data = last_model[idx] - Deltas[idx] / (hs + 1e-10)
+            new_state_dict = {
+                name: last_model[idx] - Deltas[idx] / (hs + 1e-10) for idx, name in enumerate(target_model.state_dict().keys())
+            }
+            target_model.load_state_dict(new_state_dict)
 
         else:
             # The default optimizer, FedAvg, has been applied in aggregator.py on the fly