set gpu options under hvd scenes

alibaba · Feb 22, 2024 · 8f828df · 8f828df
1 parent 47940ea
commit 8f828df
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/easy_rec/python/main.py b/easy_rec/python/main.py
@@ -105,12 +105,15 @@ def _create_estimator(pipeline_config, distribution=None, params={}):
   train_config = pipeline_config.train_config
   gpu_options = GPUOptions(allow_growth=True)  # False)
 
+  logging.info('train_config.train_distribute=%s[value=%d]' %
+     (DistributionStrategy.Name(pipeline_config.train_config.train_distribute),
+      pipeline_config.train_config.train_distribute))
+
+  # set gpu options only under hvd scenes
   if hvd is not None and pipeline_config.train_config.train_distribute in [
       DistributionStrategy.EmbeddingParallelStrategy,
       DistributionStrategy.SokStrategy, DistributionStrategy.HorovodStrategy
   ]:
-    logging.info('train_config.train_distribute=%s' %
-                 pipeline_config.train_config.train_distribute)
     local_rnk = hvd.local_rank()
     gpus = tf.config.experimental.list_physical_devices('GPU')
     logging.info('local_rnk=%d num_gpus=%d' % (local_rnk, len(gpus)))