[FIX](tf) 2D Argmax function tensor index

wbenbihi · Aug 24, 2022 · 93c9c44 · 93c9c44
1 parent 62498ab
commit 93c9c44
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/hourglass_tensorflow/utils/tf.py b/hourglass_tensorflow/utils/tf.py
@@ -212,8 +212,8 @@ def tf_matrix_argmax(tensor: tf.Tensor) -> tf.Tensor:
     """
     flat_tensor = tf.reshape(tensor, (-1, tf.shape(tensor)[-1]))
     argmax = tf.cast(tf.argmax(flat_tensor, axis=0), tf.int32)
-    argmax_x = argmax // tf.shape(tensor)[2]
-    argmax_y = argmax % tf.shape(tensor)[2]
+    argmax_x = argmax // tf.shape(tensor)[1]
+    argmax_y = argmax % tf.shape(tensor)[1]
     # stack and return 2D coordinates
     return tf.transpose(tf.stack((argmax_x, argmax_y), axis=0), [1, 0])