I'm trying to select elements differents from zero and work with them later. My input tensor has batch dimension, so I want to keep it and don't mix data over batches. I think `tf.gather_nd()` would work for me, but first I have to get the indexes of the desired data and I found `tf.where()`. I have tried the following:

```py
img = tf.constant([[[1., 0., 0.], 
                    [0., 0., 2.],
                    [0., 3, 0.]], 
                   [[1., 2., 3.], 
                    [0., 0., 1.], 
                    [0., 0., 0.]]], dtype='float32') # shape [2, 3, 3]

indexes = tf.where(tf.not_equal(img, 0.))

```

I would expect `indexes` to keep batch dimension, however it has shape `[7, 2]`. I suspect the problem comes from having different number of points that satisfies the condition in different batches. 

Is there a way to get the indexes keeping batch dimension? Thanks in advance.

**EDIT:** `indexes` has shape `[7, 3]` where first dim refers to number of points and the second dim refers to the position of the point (incluiding which batch it belongs to). But I need `indexes` to have the specific batch dimension, because later I want to use it to ghater data from `img`:

```py
Y = tf.gather_nd(img, indexes)

```

I want `Y` to have batch dimension, but as `indexes` hasn't, I get a flat tensor with data from different bateches mixed.