In fact, I think it’s pretty common use torch.max in loss function.
Like this: https://github.com/mattmacy/torchbiomed/blob/master/torchbiomed/loss.py#L32 40
If the input size is Batch x ClassNumber x H x W while the target size is Batch x H x W, many applications would use torch.max to find the right label and compare with the target.
Qi