| 233 | } |
| 234 | |
| 235 | Status OpKernelConstruction::allocate_temp(DataType type, |
| 236 | const TensorShape& shape, |
| 237 | Tensor* out_temp) { |
| 238 | AllocationAttributes attr; |
| 239 | attr.allocation_will_be_logged = true; |
| 240 | Tensor new_temp(allocator_, type, shape, attr); |
| 241 | |
| 242 | if (!new_temp.IsInitialized()) { |
| 243 | return errors::ResourceExhausted( |
| 244 | "OOM when allocating temporary tensor with shape", shape.DebugString()); |
| 245 | } |
| 246 | if (LogMemory::IsEnabled()) { |
| 247 | LogMemory::RecordTensorAllocation( |
| 248 | def_->name(), LogMemory::OP_KERNEL_CONSTRUCTION_STEP_ID, new_temp); |
| 249 | } |
| 250 | *out_temp = new_temp; |
| 251 | return Status::OK(); |
| 252 | } |
| 253 | |
| 254 | Status OpKernelConstruction::allocate_persistent( |
| 255 | DataType type, const TensorShape& shape, PersistentTensor* out_persistent, |