Hello ptrblck:
Thank you for your kind reply on above question.
I think I understand what you say, you want change the predict fucntion by adding argmax() so that when I run predict function I get the wanted shape.
Except for the method you proposed，can I try this method : I train the model without using argmax() in my forward propagation function( cause it is not differentiable),  when I  use the weight file to do an inference, I modify the model( like adding the argmax()  to forward propagation function) so that i can get the result with shape [batch_size, width, height], but I am not sure whether it works.
And also I have another question, if I modify the model( adding argmax() to model) when I implemt inference, i think there may be some probelems when I load the weight file(error may be like :   incompatible keyword etc.)