From ef09ec4d4624f1a0c792b07397c7ac703614b390 Mon Sep 17 00:00:00 2001 From: hegongshan Date: Tue, 9 Sep 2025 20:20:49 +0800 Subject: [PATCH 1/2] fix: fix num_batches --- torchrec/datasets/random.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchrec/datasets/random.py b/torchrec/datasets/random.py index 9008622e5..8acac8cec 100644 --- a/torchrec/datasets/random.py +++ b/torchrec/datasets/random.py @@ -212,7 +212,7 @@ def __init__( num_generated_batches=num_generated_batches, min_ids_per_features=min_ids_per_features, ) - self.num_batches: int = cast(int, num_batches if not None else sys.maxsize) + self.num_batches: int = cast(int, num_batches if num_batches is not None else sys.maxsize) def __iter__(self) -> Iterator[Batch]: return itertools.islice(iter(self.batch_generator), self.num_batches) From dccd743f79cd3874d736f6990c3df8f6bc50d532 Mon Sep 17 00:00:00 2001 From: hegongshan Date: Tue, 9 Sep 2025 12:49:10 +0000 Subject: [PATCH 2/2] style: format code --- torchrec/datasets/random.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchrec/datasets/random.py b/torchrec/datasets/random.py index 8acac8cec..80812705e 100644 --- a/torchrec/datasets/random.py +++ b/torchrec/datasets/random.py @@ -212,7 +212,9 @@ def __init__( num_generated_batches=num_generated_batches, min_ids_per_features=min_ids_per_features, ) - self.num_batches: int = cast(int, num_batches if num_batches is not None else sys.maxsize) + self.num_batches: int = cast( + int, num_batches if num_batches is not None else sys.maxsize + ) def __iter__(self) -> Iterator[Batch]: return itertools.islice(iter(self.batch_generator), self.num_batches)