Argmax axis not working as expected


When i define a network with tensorflow and have either
or tf.argmax(axis=3) (Fully convolutional network)

I get the wrong behavior in TensorRT, where it seems to take axis 1.
Conversion is done from frozen pb to uff with TensorRT utils and then loaded in TensorRT with Uff Parser.To get the correct behavior i have to to
tf.argmax(axis=1) (which is the image height originally)

Weirdly when I design the same network in Keras and again add tf.argmax i have to say (axis=2) to get to correct behavior in TensorRT.

I also tried to define the network Channel first and see if the error still exists but I never got any channel first network to run in TensorRT. It seems to always mess up the kernel dimensions on creation.

TRT version is 5.1

To help us debug, can you share a small repro {model or minimal repro script}? You can dm me with google drive, or dropbox links.

Hey andre,

I have encountered the exact same problem. Have you managed to find a solution to it ?