title: PyTorch how to do gathers over multiple dimensions
body: <p>I'm trying to find a way to do this without for loops.</p>
  <p>Say I have a multi-dimensional tensor <code>t0</code>:</p>
  <pre><code>bs = 4
  seq = 10
  v = 16
  t0 = torch.rand((bs, seq, v))
  </code></pre>
  <p>This has shape: <code>torch.Size([4, 10, 16])</code></p>
  <p>I have another tensor <code>labels</code> that is a batch of 5 random indices in the <code>seq</code> dimension:</p>
  <pre><code>labels = torch.randint(0, seq, size=[bs, sample])
  </code></pre>
  <p>So this has shape <code>torch.Size([4, 5])</code>. This is used to index the <code>seq</code> dimension of <code>t0</code>.</p>
  <p>What I want to do is loop over the batch dimension doing gathers using <code>labels</code> tensor. My brute force solution is this:</p>
  <pre><code>t1 = torch.empty((bs, sample, v))
  for b in range(bs):
      for idx0, idx1 in enumerate(labels[b]):
          t1[b, idx0, :] = t0[b, idx1, :]
  </code></pre>
  <p>Resulting in tensor <code>t1</code> which has shape: <code>torch.Size([4, 5, 16])</code></p>
  <p>Is there a more idiomatic way of doing this in pytorch?</p>