| 709 | } |
| 710 | |
| 711 | Status OpKernelContext::allocate_tensor( |
| 712 | DataType type, const TensorShape& shape, Tensor* out_tensor, |
| 713 | AllocatorAttributes attr, const AllocationAttributes& allocation_attr) { |
| 714 | Allocator* a = get_allocator(attr); |
| 715 | Tensor new_tensor(a, type, shape, |
| 716 | AllocationAttributes(allocation_attr.no_retry_on_failure, |
| 717 | /* allocation_will_be_logged= */ true, |
| 718 | allocation_attr.freed_by_func)); |
| 719 | |
| 720 | if (!new_tensor.IsInitialized()) { |
| 721 | return errors::ResourceExhausted( |
| 722 | "OOM when allocating tensor with shape", shape.DebugString(), |
| 723 | " and type ", DataTypeString(type), " on ", params_->device->name(), |
| 724 | " by allocator ", a->Name()); |
| 725 | } |
| 726 | if (params_->log_memory) { |
| 727 | LogMemory::RecordTensorAllocation(params_->op_kernel->name(), |
| 728 | params_->step_id, new_tensor); |
| 729 | } |
| 730 | record_tensor_reference(new_tensor); |
| 731 | *out_tensor = std::move(new_tensor); |
| 732 | return Status::OK(); |
| 733 | } |
| 734 | |
| 735 | Status OpKernelContext::allocate_output(int index, const TensorShape& shape, |
| 736 | Tensor** output, |
nothing calls this directly
no test coverage detected