title: Is there a function in Pytorch to do a batch span selection?
body: <p>I have three tensors:</p>
  <ul>
  <li><code>tokens</code> (shape <code>[N, T]</code>),</li>
  <li><code>start_index</code> of shape <strong>[N]</strong> and values in range <strong>[0, T)</strong></li>
  <li><code>end_index</code> of shape <strong>[N]</strong> and values in range <strong>[0, T)</strong></li>
  </ul>
  <p>I want to select spans from the batch of token sequences <code>tokens</code> using <code>start_index</code> and <code>end_index</code>.</p>
  <p>Essentially this is what I want:</p>
  <pre><code>[seq[s:e] for seq, s, e in zip(tokens, start_index, end_index)]
  </code></pre>
  <p>Is there a faster Pytorch-way to implement this?</p>