From 689d9974b95a764759be2028cff36735fde3001f Mon Sep 17 00:00:00 2001 From: Fei Hu Date: Mon, 5 Nov 2018 16:05:42 -0800 Subject: [PATCH] Change the maximum input for model::MakeParameter::parallelism --- .../kernels/data/experimental/numa_map_and_batch_dataset_op.cc | 2 +- tensorflow/core/kernels/data/map_and_batch_dataset_op.cc | 2 +- tensorflow/core/kernels/data/parallel_map_iterator.cc | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc index 1286795a0b2..068f8540230 100644 --- a/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc +++ b/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc @@ -244,7 +244,7 @@ class NumaMapAndBatchDatasetOp : public UnaryDatasetOpKernel { return model::MakeAsyncKnownRatioNode( std::move(args), dataset()->batch_size_, {model::MakeParameter("parallelism", num_parallel_calls_, /*min=*/1, - /*max=*/port::NumSchedulableCPUs())}); + /*max=*/ctx->runner_threadpool_size())}); } Status SaveInternal(IteratorStateWriter* writer) override { diff --git a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc index d7db3c4d183..31851925124 100644 --- a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc +++ b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc @@ -296,7 +296,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { return model::MakeAsyncKnownRatioNode( std::move(args), dataset()->batch_size_, {model::MakeParameter("parallelism", num_parallel_calls_, /*min=*/1, - /*max=*/port::NumSchedulableCPUs())}); + /*max=*/ctx->runner_threadpool_size())}); } Status SaveInternal(IteratorStateWriter* writer) override { diff --git a/tensorflow/core/kernels/data/parallel_map_iterator.cc b/tensorflow/core/kernels/data/parallel_map_iterator.cc index 10103230950..ec1c9238430 100644 --- a/tensorflow/core/kernels/data/parallel_map_iterator.cc +++ b/tensorflow/core/kernels/data/parallel_map_iterator.cc @@ -101,7 +101,7 @@ class ParallelMapIterator : public DatasetBaseIterator { std::move(args), /*ratio=*/1, {model::MakeParameter("parallelism", num_parallel_calls_, /*min=*/1, - /*max=*/port::NumSchedulableCPUs())}); + /*max=*/ctx->runner_threadpool_size())}); } Status SaveInternal(IteratorStateWriter* writer) override {