Hi ImgPreSng,
Thank you for your reply. But I could not figure out how nn.Softmax could work as torch.max function?
Could you please give some more details? Thanks so much.
For instance, here’s the DICE cost function I have. How can I use softmax replace the max op?

def forward(self, outputs, targets):
    smooth = .00001
    _, o = outputs.max(1)
    iflat = o.view(-1).float()
    tflat = targets.view(-1).float()
    intersection = (iflat * tflat).sum()

    return  1- ((2. * intersection + smooth) /
              (iflat.sum() + tflat.sum() + smooth))