From ed52a3cd1b8fe5696136e2090caba14e78869139 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-10-98.us-west-2.compute.internal>
Date: Mon, 18 Dec 2023 07:14:33 +0000
Subject: [PATCH] flush data.

---
 python/graphstorm/model/utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/graphstorm/model/utils.py b/python/graphstorm/model/utils.py
index aea73096f4..aab4cb5ca3 100644
--- 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):