fix sync optimizer bug (#391)

fix sync optimizer compatibility with pai-tf
alibaba · Jun 25, 2023 · 10668aa · 10668aa
1 parent 3e9e499
commit 10668aa
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/easy_rec/python/model/easy_rec_estimator.py b/easy_rec/python/model/easy_rec_estimator.py
@@ -206,16 +206,16 @@ def _train_model_fn(self, features, labels, run_config):
       logging.info('sync_replicas: num_worker_replias = %d' %
                    run_config.num_worker_replicas)
       if pai_util.is_on_pai():
-        extra_args = {
-            'sparse_accumulator_type': self.train_config.sparse_accumulator_type
-        }
+        optimizer = tf.train.SyncReplicasOptimizer(
+            optimizer,
+            replicas_to_aggregate=run_config.num_worker_replicas,
+            total_num_replicas=run_config.num_worker_replicas,
+            sparse_accumulator_type=self.train_config.sparse_accumulator_type)
       else:
-        extra_args = {}
-      optimizer = sync_replicas_optimizer.SyncReplicasOptimizer(
-          optimizer,
-          replicas_to_aggregate=run_config.num_worker_replicas,
-          total_num_replicas=run_config.num_worker_replicas,
-          **extra_args)
+        optimizer = sync_replicas_optimizer.SyncReplicasOptimizer(
+            optimizer,
+            replicas_to_aggregate=run_config.num_worker_replicas,
+            total_num_replicas=run_config.num_worker_replicas)
       hooks.append(
           optimizer.make_session_run_hook(run_config.is_chief, num_tokens=0))