Ask for simplified Implementation of tf.gather_nd in PyTorch

Hi, everyone.
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])[0]
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[0].item(), uv_j[1].item())
    delta[j] = dmap_j[uv_j[0].item(), uv_j[1].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 )[0]

So would developers could solve this issue or did I miss something?