From 080d9b0b8343d7f688a8c33a6649df3f4b2bec84 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Wed, 1 Oct 2025 10:50:14 -0700 Subject: [PATCH] Internal PiperOrigin-RevId: 813820929 --- grain/_src/python/dataset/transformations/batch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/grain/_src/python/dataset/transformations/batch.py b/grain/_src/python/dataset/transformations/batch.py index 765e83ec2..dc09814fb 100644 --- a/grain/_src/python/dataset/transformations/batch.py +++ b/grain/_src/python/dataset/transformations/batch.py @@ -78,6 +78,7 @@ def __call__(self, values: Sequence[T]) -> T: def _batch_fn(*xs: Sequence[T]) -> T: # If the thread pool is not available or the elements are not NumPy # arrays, fall back to the standard serial `np.stack` operation. + # TODO: Support parallel batch when elements are not NumPy if (self._parallel_batch_executor is None) or not isinstance( xs[0], np.ndarray ):