How to gather between 2 index value in batch?

1

Hello:
I am using Pytorch 0.3.1

My target Tensor (A) size is (Batch_size, sequence_len, node_dimension)
512x40x5

and my index is (Batch_size, 2-index (x,y))
512x2x20
( 0 ,.,.) =
0 8 16 … 21 29 37
7 15 23 … 22 30 38

( 1 ,.,.) =
16 8 0 … 29 21 37
23 15 7 … 30 22 38

My question is How can I gather between 2 index value in batch as follow?

To get the result (B):
( 0 ,.,.) =
A[:,0:7,:],
A[:,8:15,:],
…
A[:,37:38,:]
(1,.,.)

Some notes for the index and results :

  1. x < y (always) and the y-x>=1
  2. the length of one gathered result (B) is variable and the length always between in 2<= B <=sequence_len

I have tried in for-loop. The computation time is about ~3.5s for 512 batch size. I want to know if there is a better way to do this.

Thanks !