I wrote a small program to test the behavior of TF and the self-implemented gather_nd in PyTorch.
import tensorflow as tf
import torch
import numpy as np

import functools
import operator


def gather_nd(params, indices):
    '''
    the input indices must be a 2d tensor in the form of [[a,b,..,c],...], 
    which represents the location of the elements.
    '''
    max_value = functools.reduce(operator.mul, list(params.size())) - 1
    indices = indices.t().long()
    ndim = indices.size(0)
    idx = torch.zeros_like(indices[0]).long()
    m = 1
    
    for i in range(ndim)[::-1]:
        idx += indices[i] * m 
        m *= params.size(i)

    idx[idx < 0] = 0
    idx[idx > max_value] = 0
    return torch.take(params, idx)


idx = np.array([
	[0, 0, 0, 0],
	[0, 0, 1, 0],
	[0, 1, 1, 0],
	[2, 1, 1, 0],
	[1, 1, 4, 0],
])
mtx = np.reshape(range(120), (2, 3, 4, 5))
print(mtx)

tf_result = tf.gather_nd(mtx, idx)
with tf.Session() as sess:
	print('TF', sess.run(tf_result))

mtx = torch.from_numpy(mtx)
idx = torch.from_numpy(idx)
torch_result = gather_nd(mtx, idx)
print('Torch', torch_result)

And the output is shown as following:
TF [ 0  5 25  0  0]
Torch tensor([  0,   5,  25,   0, 100])

It seem that there are something wrong with the self-implemented one.