Argmax compares pixels in the same position across channels, and acquires the index of the highest channel.
This can be useful for semantic segmentation. Semantic segmentation models outputs the same width and height as the input image and creates a channel for each class you want to detect. For each channel, it outputs a high value if the pixel is likely the corresponding class for that channel. If you use ArgMax, you can plot which pixel belongs to which class.
