From 7c5f38b464c112f2228bbd931ac26f2fce374db2 Mon Sep 17 00:00:00 2001
From: Alex Barghi <105237337+alexbarghi-nv@users.noreply.github.com>
Date: Wed, 6 Sep 2023 13:44:33 -0400
Subject: [PATCH] [IMP] Add ability to get batch size from the loader in
 cuGraph-PyG (#3846)

Add a property getter for batch size.  Requested by JoC.

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

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)

URL: https://github.com/rapidsai/cugraph/pull/3846
---
 python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py b/python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py
index fcec341d1db..8d79685965f 100644
--- a/python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py
+++ b/python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py
@@ -351,6 +351,10 @@ def __init__(
         self.__input_nodes = input_nodes
         self.inner_loader_args = kwargs
 
+    @property
+    def batch_size(self) -> int:
+        return self.__batch_size
+
     def __iter__(self):
         self.current_loader = EXPERIMENTAL__BulkSampleLoader(
             self.__feature_store,