Hi, Thanks for your reply! I want to change the following two lines into pytorch
new_neighbors_indices = tf.batch_gather(neighbors_indices, new_neigh_inds)
new_sq_distances = tf.batch_gather(sq_distances, new_neighb_inds)

Here the neighbors_indices is dimension of [a, b] and new_neigh_inds is dimension of [a, c], I find that for the first line I can use torch.gather as,
new_neighbors_indices = torch.gather(neighbors_indices, dim=1, index=new_neighb_inds)

But the second line sq_distances is dimension of [a, b, k] and  new_neigh_inds is dimension of [a, c],  the expected results is [a, c, k] ,so how could I use torch.gather? I tried torch.gather(sq_distances, dim=1, index=new_neighb_inds) but get the error
*** RuntimeError: invalid argument 4: Index tensor must have same dimensions as input tensor at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/generic/THCTensorScatterGather.cu:16