I would like to grab tensor
v1's third dim according to indices in tensor
indices. I tried
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
incides dimensions don't match. Out should be of shape
[1,35, 20, 200]. How would I do that in PyTorch?