experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[tf.data] Fix forward compatibility issue in rebatch.cc

Browse Source 
PiperOrigin-RevId: 258820315
This commit is contained in:
Rachel Lim 2019-07-18 12:17:43 -07:00 committed by TensorFlower Gardener
parent 7dfc439541
commit fe44fdddbc
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/grappler/optimizers/data/rebatch.cc
View File

@ -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;
}