diff options
| author |  Yuefeng Zhou <yuefengz@google.com> | 2018-01-17 15:58:10 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-01-17 16:05:19 -0800 |
| commit | e86fc00d097cb8cc118435ebd464497825bcad1e (patch) | |
| tree | 2ae34164ad6b3cbd7624efa35423ff66cfe17f2c /tensorflow/core/kernels/lookup_table_op.cc | |
| parent | 04121a03acb16645dc0eac327c35a501453d8e8a (diff) | |
Combine host and device memory proto fields.
PiperOrigin-RevId: 182284426
Diffstat (limited to 'tensorflow/core/kernels/lookup_table_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/lookup_table_op.cc | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/tensorflow/core/kernels/lookup_table_op.cc b/tensorflow/core/kernels/lookup_table_op.cc index 418d9dcc61..e3872fee0e 100644 --- a/tensorflow/core/kernels/lookup_table_op.cc +++ b/tensorflow/core/kernels/lookup_table_op.cc @@ -709,8 +709,8 @@ class LookupTableInsertOp : public OpKernel { } OP_REQUIRES_OK(ctx, table->Insert(ctx, keys, values)); if (ctx->track_allocations()) { - ctx->record_host_persistent_memory_allocation(table->MemoryUsed() - - memory_used_before); + ctx->record_persistent_memory_allocation(table->MemoryUsed() - + memory_used_before); } } }; @@ -786,8 +786,8 @@ class LookupTableImportOp : public OpKernel { } OP_REQUIRES_OK(ctx, table->ImportValues(ctx, keys, values)); if (ctx->track_allocations()) { - ctx->record_host_persistent_memory_allocation(table->MemoryUsed() - - memory_used_before); + ctx->record_persistent_memory_allocation(table->MemoryUsed() - + memory_used_before); } } }; |