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);
 }