From fe44fdddbcdc16a32216b771fa5b408b89545c14 Mon Sep 17 00:00:00 2001 From: Rachel Lim Date: Thu, 18 Jul 2019 12:17:43 -0700 Subject: [PATCH] [tf.data] Fix forward compatibility issue in rebatch.cc PiperOrigin-RevId: 258820315 --- tensorflow/core/grappler/optimizers/data/rebatch.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/grappler/optimizers/data/rebatch.cc b/tensorflow/core/grappler/optimizers/data/rebatch.cc index e053405e309..b3e7f8febe3 100644 --- a/tensorflow/core/grappler/optimizers/data/rebatch.cc +++ b/tensorflow/core/grappler/optimizers/data/rebatch.cc @@ -152,7 +152,8 @@ Status MutateBatchSize(const NodeDef& node, int64 num_workers, // For all the batching datasets the batch_size is input number 1 except for // MapAndBatchDataset. int64 batch_size_arg_index = 1; - if (node.op() == "ExperimentalMapAndBatchDataset") { + if (node.op() == "ExperimentalMapAndBatchDataset" || + node.op() == "MapAndBatchDataset") { // For MapAndBatch we take the 3rd last input. batch_size_arg_index = node.input_size() - 3; }