Skip to content

Commit

Permalink
[IMP] Add ability to get batch size from the loader in cuGraph-PyG (r…
Browse files Browse the repository at this point in the history
…apidsai#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: rapidsai#3846
  • Loading branch information
alexbarghi-nv authored and rlratzel committed Sep 8, 2023
1 parent ccc8653 commit 7c5f38b
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 7c5f38b

Please sign in to comment.