Thanks KFrank,
Maybe I wasn’t that clear.
I’m trainin on loss with MSELoss between the last layer’s output and labels.
I want to apply after the last layer something  like the MaxT1 function I made in order to keep the max value of each row.
Something like:
Reshaped output from the last FCL
t1=tensor([[3.6600, 2.0000, 4.0900],
        [5.4400, 6.0000, 5.9000]])

Label
l1=tensor([[0, 0, 1],
        [0, 0, 1]])

I’d like to get from the first tensor
t2=tensor([[0, 0, 1],
        [0, 1, 0]])

and then do MSELoss(l1,t2) instead of MSELoss(l1,t1) as I’m doing now, which way higher. Also buy doing this it would be easier to check how many preds are correct.