diff --git a/python/graphstorm/gpartition/dist_partition_graph.py b/python/graphstorm/gpartition/dist_partition_graph.py
index 594774fdd..e8bdc4936 100644
--- a/python/graphstorm/gpartition/dist_partition_graph.py
+++ b/python/graphstorm/gpartition/dist_partition_graph.py
@@ -157,7 +157,8 @@ def main():
             os.path.join(output_path, "dist_graph"),
             args.metadata_filename,
             args.dgl_tool_path,
-            args.ssh_port)
+            args.ssh_port,
+            args.process_group_timeout)
 
         logging.info("DGL graph building took %f sec", dgl_graph_start - time.time())
 
@@ -241,7 +242,8 @@ def parse_args() -> argparse.Namespace:
                            help=("Whether to convert the partitioned data to the GraphBolt format "
                                "after creating the DistDGL graph."))
     argparser.add_argument("--process-group-timeout", type=int, default=1800,
-                           help="Timeout[seconds] for operations executed against the process group.")
+                           help="Timeout[seconds] for operations executed "
+                                "against the process group.")
 
     return argparser.parse_args()