Explicitely set `AutoShardPolicy.DATA` for `TensorLike` datasets
@@ -353,6 +353,12 @@ class TensorLikeDataAdapter(DataAdapter):
dataset = dataset.map(shuffle_batch)
+ options = tf.data.Options()
+ options.experimental_distribute.auto_shard_policy = (
+ tf.data.experimental.AutoShardPolicy.DATA
+ )
+ dataset = dataset.with_options(options)
+
self._dataset = dataset
def slice_inputs(self, indices_dataset, inputs):
