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

[tf.data] Fix a bug in prefetch dataset serialization logic.

Browse Source 
PiperOrigin-RevId: 313453820
Change-Id: I573d4288fbb10b7491778ce4edf24241f5e35fa1
This commit is contained in:
Jiri Simsa 2020-05-27 13:33:45 -07:00 committed by TensorFlower Gardener
parent 49a5378fb1
commit f6bf10607f
1 changed files with 7 additions and 3 deletions

10
tensorflow/core/kernels/data/prefetch_dataset_op.cc
View File

@ -100,9 +100,13 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
TF_RETURN_IF_ERROR(b->AddScalar(buffer_size_, &buffer_size));
AttrValue slack_period_attr;
b->BuildAttrValue(slack_period_, &slack_period_attr);
TF_RETURN_IF_ERROR(b->AddDataset(
this, {input_graph_node, buffer_size},
{std::make_pair(kSlackPeriod, slack_period_attr)}, output));
AttrValue legacy_autotune_attr;
b->BuildAttrValue(legacy_autotune_, &legacy_autotune_attr);
TF_RETURN_IF_ERROR(
b->AddDataset(this, {input_graph_node, buffer_size},
{std::make_pair(kSlackPeriod, slack_period_attr),
std::make_pair(kLegacyAutotune, legacy_autotune_attr)},
output));
return Status::OK();
}