1
$\begingroup$

I would like to grab tensor v1's third dim according to indices in tensor indices. I tried torch.take and torch.gather, but no go. See example:

import torch
v1 = torch.randn((1,35,20,200))
indices = torch.randint(0, 20, (1,35,20))
out = torch.gather(v1, 2, indices)

This gives an error as v1 and incides dimensions don't match. Out should be of shape [1,35, 20, 200]. How would I do that in PyTorch?

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.