diff --git a/tensorflow/core/kernels/data/dataset_test_base.cc b/tensorflow/core/kernels/data/dataset_test_base.cc index b03abed9359..54de6888508 100644 --- a/tensorflow/core/kernels/data/dataset_test_base.cc +++ b/tensorflow/core/kernels/data/dataset_test_base.cc @@ -190,6 +190,7 @@ Status DatasetOpsTestBase::CreateIteratorContext( OpKernelContext* const op_context, std::unique_ptr* iterator_context) { IteratorContext::Params params(op_context); + params.resource_mgr = op_context->resource_manager(); function_handle_cache_ = absl::make_unique(flr_); params.function_handle_cache = function_handle_cache_.get(); *iterator_context = absl::make_unique(params); @@ -228,6 +229,7 @@ Status DatasetOpsTestBase::InitFunctionLibraryRuntime( TF_RETURN_IF_ERROR(DeviceFactory::AddDevices( options, "/job:localhost/replica:0/task:0", &devices)); device_mgr_ = absl::make_unique(std::move(devices)); + resource_mgr_ = absl::make_unique("default_container"); FunctionDefLibrary proto; for (const auto& fdef : flib) *(proto.add_function()) = fdef; @@ -269,6 +271,7 @@ Status DatasetOpsTestBase::CreateOpKernelContext( step_container_ = absl::make_unique(0, [](const string&) {}); params_->step_container = step_container_.get(); + params_->resource_manager = resource_mgr_.get(); checkpoint::TensorSliceReaderCacheWrapper slice_reader_cache_wrapper; slice_reader_cache_ = absl::make_unique(); diff --git a/tensorflow/core/kernels/data/dataset_test_base.h b/tensorflow/core/kernels/data/dataset_test_base.h index f3a0d0ae12d..ca2be6b9258 100644 --- a/tensorflow/core/kernels/data/dataset_test_base.h +++ b/tensorflow/core/kernels/data/dataset_test_base.h @@ -206,6 +206,7 @@ class DatasetOpsTestBase : public ::testing::Test { std::function)> runner_; std::unique_ptr device_mgr_; std::unique_ptr lib_def_; + std::unique_ptr resource_mgr_; std::unique_ptr params_; std::unique_ptr slice_reader_cache_;