[FIX] Improve the Warning Displayed for Uneven Batch Distributions (#…

…4920)

In order to better debug issues with batch distribution, this PR updates the warning message to show how many batches each rank received.  Partially resolves rapidsai/cugraph-gnn#130.

Authors:
  - Alex Barghi (https://github.com/alexbarghi-nv)

Approvers:
  - Brad Rees (https://github.com/BradReesWork)

URL: #4920

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)