1 files changed, 11 insertions, 2 deletions
diff --git a/tensorflow/core/kernels/lookup_table_op.h b/tensorflow/core/kernels/lookup_table_op.h
index 4cd25a3cc6..ff23a09a24 100644
--- a/tensorflow/core/kernels/lookup_table_op.h
+++ b/tensorflow/core/kernels/lookup_table_op.h
@@ -64,8 +64,8 @@ class LookupTableOp : public OpKernel {
         return ctx->status();
       }
       if (ctx->track_allocations()) {
-        ctx->record_device_persistent_memory_allocation(
-            container->MemoryUsed());
+        ctx->record_host_persistent_memory_allocation(
+            container->MemoryUsed() + table_handle_.AllocatedBytes());
       }
       *ret = container;
       return Status::OK();
@@ -225,6 +225,15 @@ class HashTable : public InitializableLookupTable {
     return Status::OK();
   }
 
+  int64 MemoryUsed() const override {
+    if (table_) {
+      const int64 num_elements = table_->size();
+      return num_elements * (sizeof(K) + sizeof(V));
+    } else {
+      return 0;
+    }
+  }
+
  private:
   std::unique_ptr<std::unordered_map<K, V>> table_;
 };