There have been some discussions about the topic, but none of them solved my problem.
Existing discussion threads:

How to use Tf.gather_nd in pytorch 34
How to do the tf.gather_nd in pytorch? 23
How to implement an equivalent of tf.gather in pytorch 23

I am not familiar with TF. And I want to translate a TF code to PyTorch.
The TF code goes as follows:
sample_idx = tf.cast(tf.round(sample_idx), 'int32')
g_val = tf.gather_nd(sample_grid, sample_idx) 

where sample_idx is of size [3211264, 4] and sample_grid is of size [1, 1, 32, 32, 32].
Moreover, there are negative values in the sample_idx. And I tried to print it as follows:
tensor([[  0., -18.,  16.,  21.],
        [  0., -18.,  16.,  20.],
        [  0., -18.,  16.,  20.],
        ...,
        [  0.,  56.,  -0.,  -5.],
        [  0.,  56.,  -0.,  -5.],
        [  0.,  56.,  -0.,  -5.]], device='cuda:0')

And the g_val is of size [3211264] in TF.
I wonder how to do this operation with PyTorch?
Thanks for your help.