diff options
Diffstat (limited to 'tensorflow/core/common_runtime/eager/context.h')
-rw-r--r-- | tensorflow/core/common_runtime/eager/context.h | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h index eb6eb0d55a..5ed6057ec6 100644 --- a/tensorflow/core/common_runtime/eager/context.h +++ b/tensorflow/core/common_runtime/eager/context.h @@ -33,6 +33,7 @@ limitations under the License. #include "tensorflow/core/distributed_runtime/eager/eager_client.h" #include "tensorflow/core/distributed_runtime/server_lib.h" #endif +#include "tensorflow/core/framework/log_memory.h" #include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/core/threadpool.h" @@ -141,6 +142,7 @@ class EagerContext { void AddKernelToCache(Fprint128 cache_key, KernelAndDevice* kernel); bool LogDevicePlacement() { return log_device_placement_; } + bool LogMemory() { return log_memory_; } Rendezvous* GetRendezvous() { return rendezvous_; } @@ -261,6 +263,8 @@ class EagerContext { std::unordered_map<std::thread::id, bool> thread_local_async_ GUARDED_BY(async_map_mu_); + const bool log_memory_; + Env* const env_; #ifndef __ANDROID__ |