cannot compute argmax and copy back to cpu when use Tensorrt

[hlyf-xs]

Hello, This is a great job! But I found a question that the program cannot compute argmax and copy back to cpu when I use tensorrt to speed up the model. So It's very slow beacuse use np.argmax()

[MagicHly]

Hello,I also found that it is too slow to use np.argmax() recently.How can I sovle this question?