[fix new benchmark]

BamLubi · Sep 21, 2022 · 5cbb464 · 5cbb464
1 parent 59feef2
commit 5cbb464
Showing 1 changed file with 10 additions and 12 deletions.
diff --git a/tools/static_ps_trainer.py b/tools/static_ps_trainer.py
@@ -202,18 +202,16 @@ def run_worker(self):
             self.train_result_dict["speed"].append(epoch_speed)
 
             model_dir = "{}/{}".format(save_model_path, epoch)
-            if fleet.is_first_worker() and save_model_path:
-                if is_distributed_env():
-                    fleet.save_inference_model(
-                        self.exe, model_dir,
-                        [feed.name for feed in self.inference_feed_var],
-                        self.inference_target_var)
-                else:
-                    paddle.static.save_inference_model(
-                        model_dir,
-                        [feed.name for feed in self.inference_feed_var],
-                        [self.inference_target_var], self.exe)
-            fleet.barrier_worker()
+            if is_distributed_env():
+                fleet.save_inference_model(
+                    self.exe, model_dir,
+                    [feed.name for feed in self.inference_feed_var],
+                    self.inference_target_var)
+            else:
+                paddle.static.save_inference_model(
+                    model_dir,
+                    [feed.name for feed in self.inference_feed_var],
+                    [self.inference_target_var], self.exe)
 
         if reader_type == "InmemoryDataset":
             self.reader.release_memory()