Slicing a 3D tensor using two index tensor

vision
#1

Hi,
I have one 3D tensor X of shape [5, 16, 128] and I also have two 2-D tensors P and N containing indices of X in some order. P is of shape [5,6] and N is of shape [5,10].

How do I slice X using P and N to get two tensors of shape X1 = [5,6,128] and X2 = [5,10,128]?

X = torch.randn(5, 16, 128)
P = torch.tensor([[0, 1, 2, 3, 4, 5],
        [0, 1, 2, 3, 4, 5],
        [0, 1, 2, 3, 4, 5],
        [0, 1, 2, 3, 4, 5],
        [0, 1, 2, 3, 4, 5]])
N = torch.tensor([[ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
        [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
        [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
        [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
        [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15]])

X[P, :] = ? 
X[N, :] = ?
*** IndexError: index 5 is out of bounds for dimension 0 with size 5

I tried X[P, :] since I thought the 3rd dim (of shape 128) would get carried forward with : but I get an error. How do I do this in pytorch? (in Tensorflow i am aware of tf.gather)

thanks,
srk

Note: first first dim of index tensors P and N is same as X and second dim of P and N (6+10) add to the second dim of X (16).

#2

If you want to simply split X into two tensors of shape [5,6,128] and [5,10,128] then you can use

X1,X2 = torch.split(X,[6,10])