hao.liu:

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.


Yes, this would be my proposed approach if you can’t add the argmax operation in your deployment setup. Manipulating the forward would work or probably the cleaner approach would be to write a custom predict method which adds the argmax operation to it as described before. You could alsoreturn two outputs in the forward: the logits as well as the class predictions and could use one of them corresponding to training or inference. The “best” approach depends on your workflow.



 hao.liu:

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.)


No, you should not see any state_dict mismatches, as torch.argmax is a purely functional call, won’t have any parameters, and will thus not be stored in the state_dict.