diff options
| author |  Yuefeng Zhou <yuefengz@google.com> | 2017-01-20 10:59:37 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-01-20 11:06:19 -0800 |
| commit | 448c0dc61c40a631bf5f1a9ec1327591a28e63e6 (patch) | |
| tree | c1aaa2cf06913f6e1c653898146b8ca555263203 /tensorflow/core/framework/allocator.cc | |
| parent | 8499dc1ee56b40a6e1f16ccbf3073a12523a0f58 (diff) | |
Fix the bug that when cpu_allocated is called before session initialization memory cannot be tracked.
Change: 145103741
Diffstat (limited to 'tensorflow/core/framework/allocator.cc')
| -rw-r--r-- | tensorflow/core/framework/allocator.cc | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/core/framework/allocator.cc b/tensorflow/core/framework/allocator.cc index 812ce4bfe7..ff31ad965b 100644 --- a/tensorflow/core/framework/allocator.cc +++ b/tensorflow/core/framework/allocator.cc @@ -120,6 +120,9 @@ Allocator* MakeCpuAllocator() { Allocator* cpu_allocator() { static Allocator* cpu_alloc = MakeCpuAllocator(); + if (cpu_allocator_collect_full_stats && !cpu_alloc->TracksAllocationSizes()) { + cpu_alloc = new TrackingAllocator(cpu_alloc, true); + } return cpu_alloc; } |