| 98 | } |
| 99 | |
| 100 | Allocator* ThreadPoolDevice::GetScopedAllocator(AllocatorAttributes attr, |
| 101 | int64 step_id) { |
| 102 | if (attr.scope_id > 0) { |
| 103 | return scoped_allocator_mgr_->GetContainer(step_id)->GetInstance( |
| 104 | attr.scope_id); |
| 105 | } |
| 106 | LOG(FATAL) << "Unexpected call to ThreadPoolDevice::GetScopedAllocator " |
| 107 | << "attr.scope_id = " << attr.scope_id; |
| 108 | return allocator_; |
| 109 | } |
| 110 | |
| 111 | Status ThreadPoolDevice::MakeTensorFromProto( |
| 112 | const TensorProto& tensor_proto, const AllocatorAttributes alloc_attrs, |