title: how to use scipy.optimize.linear_sum_assignment in tensorflow or keras?
body: <p>first time posting here ! If my question is lacking anything please tell me and I'll fix it !</p>
  
  <p>Facebook recently released DETR, an object detection model using transformers ! The model is implemented with Pytorch and I'm trying to implement the loss function where Hungarian algorithm is involved but with Keras and Tensorflow as a custom loss function for Keras model. In the original implementation from Facebook, it's line 81-82 in <a href="https://github.com/facebookresearch/detr/blob/master/models/matcher.py" rel="noreferrer">https://github.com/facebookresearch/detr/blob/master/models/matcher.py</a> </p>
  
  <p>In order to use numpy and classic python function, I used:</p>
  
  <pre><code>    def hungarian_loss(losses):
          row_ind, col_ind = linear_sum_assignment(losses)
          idx = [[i, j] for i, j in zip(row_ind, col_ind)]
          return idx
  
      # dist loss is a 5x5 matrix, and idx is 5x2 indexes
      idx = tf.py_function(func=hungarian_loss, inp=[dist_loss], Tout=tf.int32)
      min_val = tf.gather_nd(dist_loss, idx)
      return K.mean(min_val)
  </code></pre>
  
  <p>But I got :</p>
  
  <pre><code>tensorflow.python.framework.errors_impl.InvalidArgumentError:  Inner dimensions of output shape must match inner dimensions of updates shape. Output: [5,5] updates: [5]
  </code></pre>
  
  <p>Is it because I'm trying to use something that wasn't a tf.Tensor as loss ? </p>