From ea620701abe69f66b1cc4e0ea772eea21ae2e51e Mon Sep 17 00:00:00 2001 From: Xiang Song Date: Fri, 6 Oct 2023 23:58:31 -0700 Subject: [PATCH] update --- python/graphstorm/gconstruct/remap_result.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/graphstorm/gconstruct/remap_result.py b/python/graphstorm/gconstruct/remap_result.py index 268d8f2c4f..5b8c85e8e0 100644 --- a/python/graphstorm/gconstruct/remap_result.py +++ b/python/graphstorm/gconstruct/remap_result.py @@ -58,8 +58,8 @@ def worker_remap_edge_pred(pred_file_path, src_nid_path, for i in range(num_chunks): output_fname = f"{output_fname_prefix}_{pad_file_index(i)}.parquet" - start = i * num_chunks - end = (i + 1) * num_chunks if i + 1 < num_chunks else len(pred_result) + start = i * chunk_size + end = (i + 1) * chunk_size if i + 1 < chunk_size else len(pred_result) pred = pred_result[start:end] src_nid = src_id_map.map_id(src_nids[start:end]) dst_nid = dst_id_map.map_id(dst_nids[start:end])