I have a high dimensional tensor with shape x
and the index i want to take values idx:
B, N = 2, 10
x = torch.zeros(B, 1, 128, 128)
idx = torch.randint(128, size=(B, N, 2))`
where idx[:, :, 0], idx[:, :, 1] are where i want to take values from x.
Also, I want to specify a radius around the index I want to take values.
For instance, radius = 6
I tried to do
i = idx[:, :, 0]
j = idx[:, :, 1]
x[:, :, i-radius:i+radius, j-radius:j+radius]
but got TypeError: only integer tensors of a single element can be converted to an index.
Is there any way I can do this without looping over the first and second dimension of idx?
Thanks!