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?


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.