flush data.

awslabs · Dec 18, 2023 · ed52a3c · ed52a3c
1 parent 03b6002
commit ed52a3c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/graphstorm/model/utils.py b/python/graphstorm/model/utils.py
@@ -32,6 +32,7 @@
 from ..gconstruct.file_io import stream_dist_tensors_to_hdf5
 from ..utils import get_rank, barrier, get_world_size, create_dist_tensor
 from ..data.utils import alltoallv_cpu, alltoallv_nccl
+from ..distributed import flush_data
 
 # placeholder of the ntype for homogeneous graphs
 NTYPE = dgl.NTYPE
@@ -1039,7 +1040,7 @@ def _load_id_mapping(self, g, ntype, id_mappings):
                 f"Expect {id_mapping.shape[0]}, but get {num_nodes}"
             # Save ID mapping into dist tensor
             id_mapping_info[th.arange(num_nodes)] = id_mapping
-        barrier()
+        flush_data()
         return id_mapping_info
 
     def shuffle_nids(self, ntype, nids):