Update flow.py

Signed-off-by: Parth Mandaliya <parthx.mandaliya@intel.com>
ParthMandaliya · Sep 28, 2023 · 1c53e30 · 1c53e30
1 parent 07f283f
commit 1c53e30
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/openfl-workspace/experimental/501_pytorch_tinyimagenet_transfer_learning/src/flow.py b/openfl-workspace/experimental/501_pytorch_tinyimagenet_transfer_learning/src/flow.py
@@ -344,6 +344,7 @@ def inference(network, test_loader):
 
 
 def fedavg(models):
+    models = [model.to('cpu') for model in models]
     new_model = models[0]
     state_dicts = [model.state_dict() for model in models]
     state_dict = new_model.state_dict()