I found that the results are not the same as TF when batch size is larger than 1.
For example, the size of sample_idx and sample_grid are [57802752, 4] and [6, 32, 32, 32], respectively.
sample_grid = torch.from_numpy(np.load('sample_grid.npy', sample_grid)).cuda()
sample_idx = torch.from_numpy(np.load('sample_idx.npy', sample_idx)).cuda()

g_val = gather_nd(sample_grid, sample_idx)
print(torch.sum(g_val[0, 0, 0, ...]))   # tensor(52197., device='cuda:0')

However, in TF:
_sample_grid = np.load('sample_grid.npy', sample_grid)
_sample_idx = np.load('sample_idx.npy', sample_idx)

g_val = tf.gather_nd(sample_grid, sample_idx)
_g_val = sess.run(g_val, {
    sample_grid: _sample_grid,
    sample_idx: _sample_idx
})
print(np.sum(_g_val[0, 0, 0, ...]))    # 43380.0

You can download the sample_grid.npy and sample_idx.npy from here 1.
Who can help me out