How to implement tf.gather_nd with argument batch_dims in Pytorch?

vision
2

This post might provide you a matching implementation for TensorFlow’s gather_nd.
Let me know, if that works for you.