Made sure that the tracking allocator always counts the allocated sizes.

Made the corresponding unit test more robust.
Change: 115575179
This commit is contained in:
Benoit Steiner 2016-02-25 10:12:26 -08:00 committed by TensorFlower Gardener
parent 5c9f4f8973
commit 63bd3efc5c
2 changed files with 3 additions and 3 deletions

View File

@ -58,7 +58,7 @@ void* TrackingAllocator::AllocateRaw(
in_use_.emplace(std::make_pair(ptr, chunk)); in_use_.emplace(std::make_pair(ptr, chunk));
allocated_ += allocated_bytes; allocated_ += allocated_bytes;
high_watermark_ = std::max(high_watermark_, allocated_); high_watermark_ = std::max(high_watermark_, allocated_);
total_bytes_ += num_bytes; total_bytes_ += allocated_bytes;
++ref_; ++ref_;
} else { } else {
mutex_lock lock(mu_); mutex_lock lock(mu_);

View File

@ -96,8 +96,8 @@ TEST(TrackingAllocatorTest, SimpleNoTracking) {
sizes = ta->GetSizesAndUnRef(); sizes = ta->GetSizesAndUnRef();
EXPECT_EQ(16, sizes.first); EXPECT_LE(16, sizes.first);
EXPECT_EQ(12, sizes.second); EXPECT_LE(12, sizes.second);
ta->DeallocateRaw(p2); ta->DeallocateRaw(p2);
} }