I am now trying to implement tf.gather_nd using Pytorch.
I have tried several methods discussed here, but none of them could meet my need.
Here I give my code which has been proven to able to work,but I want someone could gvie me a more simplified code using pytorch.
Code are showed below
import torch uv = torch.ones(42).view([-1, 2]).long() print("uv=", uv.shape) dmap = torch.rand([1, 32, 32, 21, 3]) print("dmap.shape", dmap.shape) delta = torch.zeros((21, 3)).float().cuda() for j in range(21): uv_j = uv[j].long() dmap_j = dmap[:, :, j, :].squeeze() # print("dmap_j.shape=", dmap_j.shape) # print(uv_j.item(), uv_j.item()) delta[j] = dmap_j[uv_j.item(), uv_j.item()] print("delta.shape=", delta.shape) print("delta=", delta)
In tensorflow it is really simple:
delta = tf.gather_nd( tf.transpose(dmap, [0, 3, 1, 2, 4]), uv, batch_dims=2 )
So would developers could solve this issue or did I miss something?