Method SetMemory

tensorflow/core/common_runtime/step_stats_collector.cc:151–161

}

void NodeExecStatsWrapper::SetMemory(OpKernelContext* ctx) {
  for (const auto& allocator_pair : ctx->ConsumeWrappedAllocators()) {
    AddAllocation(allocator_pair.first, allocator_pair.second);
  }
  auto* ms = stats_->mutable_memory_stats();
  ms->set_temp_memory_size(ctx->temp_memory_allocated());
  for (const auto& alloc_id : ctx->persistent_alloc_ids()) {
    ms->mutable_persistent_tensor_alloc_ids()->Add(alloc_id);
  }
  ms->set_persistent_memory_size(ctx->persistent_memory_allocated());
}

void NodeExecStatsWrapper::SetOutput(int slot, const Tensor* tensor) {
  DCHECK(tensor);