| 393 | } |
| 394 | |
| 395 | void XRTTupleAllocation::SetDeviceMemorySize() { |
| 396 | size_t size = 0; |
| 397 | for (auto& index_buffer : buffers_) { |
| 398 | if (index_buffer.second != nullptr) { |
| 399 | size += index_buffer.second->allocation().size(); |
| 400 | } |
| 401 | } |
| 402 | device_memory_size_ = size; |
| 403 | } |
| 404 | |
| 405 | /* static */ Status XRTTupleAllocation::ExpandTreeOfTuples( |
| 406 | const xla::ShapeTree<ExpandedTupleInput>& elements, int device_ordinal, |
no test coverage detected