Hi,
How to use the chunk() function to process the following situation?
>>> w = torch.randn(3, 4, 5)
>>> w
tensor([[[-0.7057,  0.2028,  0.5647,  0.6096, -0.3521],
         [ 0.6071, -0.1356,  0.4405, -1.0726, -0.0379],
         [ 1.0409, -1.6037, -1.3454, -1.4991, -0.5502],
         [ 0.9706,  1.2075, -0.8726,  0.3872, -1.2042]],

        [[ 1.5871,  1.0355, -0.3684, -0.0823, -0.1831],
         [ 0.7262,  2.4277, -0.1173,  0.4026,  0.8243],
         [ 1.1217, -0.6436,  1.4286,  0.6216,  0.0021],
         [ 1.0578, -0.4555, -0.2754,  0.6451,  0.9100]],

        [[ 0.1190, -0.9509,  0.2875, -0.9116,  0.6941],
         [ 0.4358, -0.6536, -0.5309, -0.5608,  0.6856],
         [ 0.0518, -0.8806,  0.4239,  1.2213,  0.7129],
         [ 0.3308, -1.5465,  1.9394, -1.4332,  1.3698]]])
>>> w.shape
torch.Size([3, 4, 5]) (batch_size, time_step, hidden)

>>> index = torch.tensor([[0, 1], [1, 2], [0, 3]])  (batch_size, k)

I want to get this result when dim=1:
tensor([[[-0.7057,  0.2028,  0.5647,  0.6096, -0.3521],
         [ 0.6071, -0.1356,  0.4405, -1.0726, -0.0379]],

        [[ 0.7262,  2.4277, -0.1173,  0.4026,  0.8243],
         [ 1.1217, -0.6436,  1.4286,  0.6216,  0.0021]],

        [[ 0.1190, -0.9509,  0.2875, -0.9116,  0.6941],
         [ 0.0518, -0.8806,  0.4239,  1.2213,  0.7129]]])
shape :(3, 2, 5)

What is the most effective and simple way to do this?
Thanks!!