Using torch.nonzero to index into a larger tensor

2

Hi Dave,

You could have a try on the code snippet as follows, I’m not sure whether it is efficient, but it works.

z = x[:, idx[:, 0], idx[:, 1]]
z = torch.transpose(z, 0, 1)
1 Like