diff options
Diffstat (limited to 'tensorflow/contrib/data/ops/dataset_ops.cc')
| -rw-r--r-- | tensorflow/contrib/data/ops/dataset_ops.cc | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/contrib/data/ops/dataset_ops.cc b/tensorflow/contrib/data/ops/dataset_ops.cc index 66a7c7fdcd..cc5e250ea1 100644 --- a/tensorflow/contrib/data/ops/dataset_ops.cc +++ b/tensorflow/contrib/data/ops/dataset_ops.cc @@ -168,9 +168,11 @@ output_shapes: The list of shapes being produced. REGISTER_OP("MultiDeviceIteratorInit") .Input("dataset: variant") .Input("multi_device_iterator: resource") + .Input("max_buffer_size: int64") .Output("incarnation_id: int64") .Doc(R"doc( Initializes the multi device iterator with the given dataset. +max_buffer_size: The maximum size of the host side per device buffer to keep. incarnation_id: An int64 indicating which incarnation of the MultiDeviceIterator is running. dataset: Dataset to be iterated upon. |