Set GPU index based on local_rank instead of rank

ChenYuHo · web-flow · commit a5f26eeef7ae · 2021-08-22T17:45:54.000+08:00
diff --git a/TensorFlow/Recommendation/NCF/ncf.py b/TensorFlow/Recommendation/NCF/ncf.py
@@ -227,7 +227,7 @@ def main():
     # Create and run Data Generator in a separate thread
     data_generator = DataGenerator(
         args.seed,
-        hvd.rank(),
+        hvd.local_rank(),
         nb_users,
         nb_items,
         neg_mat,
