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 :

x < y (always) and the y-x>=1
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 !