| 88 | |
| 89 | protected: |
| 90 | dim_t get_batch_index(const dim_t batch_size, |
| 91 | const dim_t batch_id, |
| 92 | const std::vector<dim_t>& batch_offset) const { |
| 93 | const auto beam_size = batch_size / batch_offset.size(); |
| 94 | return batch_offset[batch_id / beam_size]; |
| 95 | } |
| 96 | |
| 97 | dim_t get_sample_begin(const dim_t batch_size, |
| 98 | const dim_t batch_id, |