diff --git a/tensorflow_datasets/core/writer.py b/tensorflow_datasets/core/writer.py
index 8b8f8eacaa1..0b89b40acf7 100644
--- a/tensorflow_datasets/core/writer.py
+++ b/tensorflow_datasets/core/writer.py
@@ -716,7 +716,7 @@ def write_from_pcollection(self, examples_pcollection):
         # (_ShardInfo)
         | "CollectShardInfo" >> beam.transforms.combiners.ToList()
         # [_ShardInfo]
-        | "CalculateSplitInfo" >> beam.ParDo(self._store_split_info)
+        | "CalculateSplitInfo" >> beam.Map(self._store_split_info)
     )
 
   def finalize(self) -> tuple[list[int], int]: