From 554b64f0827f29fd3f2851f645c4bd805f44e927 Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Wed, 5 Feb 2025 12:06:51 -0800 Subject: [PATCH] improve warning message --- python/cugraph/cugraph/gnn/data_loading/dist_sampler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py b/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py index 6ac4ca142ee..9646777a109 100644 --- a/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py +++ b/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py @@ -149,7 +149,8 @@ def get_start_batch_offset( warnings.warn( "Not all ranks received the same number of batches. " "This might cause your training loop to hang " - "due to uneven inputs." + "due to uneven inputs. This is the number of " + f"batches receieved on each rank: {t.tolist()}." ) return (0 if rank == 0 else t.cumsum(dim=0)[rank - 1], input_size_is_equal)