Update clientapfl.py

TsingZ0 · Apr 1, 2023 · daac33a · daac33a
1 parent a30c0cc
commit daac33a
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/system/flcore/clients/clientapfl.py b/system/flcore/clients/clientapfl.py
@@ -27,6 +27,7 @@ def train(self):
 
         # self.model.to(self.device)
         self.model.train()
+        self.model_per.train()
 
         max_local_steps = self.local_steps
         if self.train_slow:
@@ -110,9 +111,7 @@ def test_metrics(self):
 
     def train_metrics(self):
         trainloader = self.load_train_data()
-        # self.model = self.load_model('model')
-        # self.model.to(self.device)
-        self.model.eval()
+        self.model_per.train()
 
         train_num = 0
         losses = 0
@@ -128,7 +127,4 @@ def train_metrics(self):
                 train_num += y.shape[0]
                 losses += loss_per.item() * y.shape[0]
 
-        # self.model.cpu()
-        # self.save_model(self.model, 'model')
-
         return losses, train_num