From 63bd3efc5c9a052d605766497726120b73f5cb3c Mon Sep 17 00:00:00 2001
From: Benoit Steiner <benoit.steiner.goog@gmail.com>
Date: Thu, 25 Feb 2016 10:12:26 -0800
Subject: [PATCH] Made sure that the tracking allocator always counts the
 allocated sizes. Made the corresponding unit test more robust. Change:
 115575179

---
 tensorflow/core/framework/tracking_allocator.cc      | 2 +-
 tensorflow/core/framework/tracking_allocator_test.cc | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/framework/tracking_allocator.cc b/tensorflow/core/framework/tracking_allocator.cc
index 5e6fe24d5a7..62d94bdf34b 100644
--- a/tensorflow/core/framework/tracking_allocator.cc
+++ b/tensorflow/core/framework/tracking_allocator.cc
@@ -58,7 +58,7 @@ void* TrackingAllocator::AllocateRaw(
     in_use_.emplace(std::make_pair(ptr, chunk));
     allocated_ += allocated_bytes;
     high_watermark_ = std::max(high_watermark_, allocated_);
-    total_bytes_ += num_bytes;
+    total_bytes_ += allocated_bytes;
     ++ref_;
   } else {
     mutex_lock lock(mu_);
diff --git a/tensorflow/core/framework/tracking_allocator_test.cc b/tensorflow/core/framework/tracking_allocator_test.cc
index 1e68d156ffc..3c4d93a182d 100644
--- a/tensorflow/core/framework/tracking_allocator_test.cc
+++ b/tensorflow/core/framework/tracking_allocator_test.cc
@@ -96,8 +96,8 @@ TEST(TrackingAllocatorTest, SimpleNoTracking) {
 
   sizes = ta->GetSizesAndUnRef();
 
-  EXPECT_EQ(16, sizes.first);
-  EXPECT_EQ(12, sizes.second);
+  EXPECT_LE(16, sizes.first);
+  EXPECT_LE(12, sizes.second);
 
   ta->DeallocateRaw(p2);
 }