Change the order of batching to avoid None

PiperOrigin-RevId: 689075797

kauldron/data/py/base.py

@@ -84,9 +84,6 @@ def ds_with_transforms(self, rng: random.PRNGKey) -> grain.MapDataset:
 
     ds = transform_utils.apply_transforms(ds, self.transforms)
 
-    if self.batch_size:
-      ds = ds.batch(self.batch_size, drop_remainder=self.batch_drop_remainder)
-
     return ds
 
   @functools.cached_property
@@ -111,6 +108,10 @@ def _root_ds(self) -> grain.IterDataset:
     # `_root_map_ds` because `_root_map_ds` does not propagate `len`
     ds = self._root_map_ds
     ds = ds.to_iter_dataset(read_options=self.read_options)
+    # We do batching after conversion to `IterDataset` to avoid None during
+    # batching.
+    if self.batch_size:
+      ds = ds.batch(self.batch_size, drop_remainder=self.batch_drop_remainder)
 
     # Distribute the execution across multiple worker processes.
     num_workers = _get_num_workers(self.num_workers)