I preprocessed the indices before torch.take and make sure that the index value in each dimension is correct.
The modified gather_nd is listed as follows:
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.
    '''
    # Normalize indices values
    params_size = list(params.size())
    assert len(params_size) == indices.size(1)

    indices[indices < 0] = 0
    for idx, ps in enumerate(params_size):
        indices[indices[:, idx] >= ps] = 0

    # Generate indices
    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)

    return torch.take(params, idx)