| 764 | } |
| 765 | |
| 766 | void FinishTimeStep() { |
| 767 | for (const auto& p : copy_out_) { |
| 768 | const Tensor& aligned = p.second; |
| 769 | Tensor original = p.first; |
| 770 | // Copy from aligned back to original. |
| 771 | functor::TensorCopyToUnaligned<Device, T>()(device_, aligned.flat<T>(), |
| 772 | original.unaligned_flat<T>()); |
| 773 | } |
| 774 | copy_out_.clear(); |
| 775 | // Mark all entries as not in use. |
| 776 | for (auto& entry : pool_) { |
| 777 | entry.second.second = false; |
| 778 | } |
| 779 | } |
| 780 | |
| 781 | private: |
| 782 | // Return a slice at position 'pos'. Result may be unaligned. The resulting |