I liked your post Extract batchwise multiple coordinates with pytorch functions (similar to tf.gather_nd) - #5 by ptrblck.
Thanks.