diff options
Diffstat (limited to 'tensorflow/core/common_runtime/eager/kernel_and_device.h')
-rw-r--r-- | tensorflow/core/common_runtime/eager/kernel_and_device.h | 11 |
1 files changed, 6 insertions, 5 deletions
diff --git a/tensorflow/core/common_runtime/eager/kernel_and_device.h b/tensorflow/core/common_runtime/eager/kernel_and_device.h index 0ef419cbaa..04151a1171 100644 --- a/tensorflow/core/common_runtime/eager/kernel_and_device.h +++ b/tensorflow/core/common_runtime/eager/kernel_and_device.h @@ -52,12 +52,12 @@ class KernelAndDevice { static Status Init(const NodeDef& ndef, FunctionLibraryRuntime* flib, std::function<void(std::function<void()>)>* runner, KernelAndDevice* out); - // TODO(ashankar): Remove this - static Status InitOp(Device* device, const NodeDef& ndef, - KernelAndDevice* out); - KernelAndDevice(tensorflow::Rendezvous* rendez) - : device_(nullptr), flib_(nullptr), rendez_(rendez) {} + KernelAndDevice(tensorflow::Rendezvous* rendez, bool log_memory) + : device_(nullptr), + flib_(nullptr), + rendez_(rendez), + log_memory_(log_memory) {} // TODO(ashankar): Handle list-valued inputs. Status Run(std::vector<Tensor>* inputs, std::vector<Tensor>* outputs, @@ -87,6 +87,7 @@ class KernelAndDevice { DataTypeVector output_dtypes_; std::function<void(std::function<void()>)>* runner_; std::function<void(std::function<void()>)> default_runner_; + const bool log_memory_; }; } // namespace tensorflow |