Hello Amos!



 Amos_Brazzoli:

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.


Have you tried putting the stuff that isn’t part of your actual
training – the stuff “like the MaxT1 function” – inside of a
with torch.no_grad(): block?  You want to keep
calculations that don’t contribute to the loss you are training
on out of the gradients used for training so they don’t pollute
things or cause errors.  That is the purpose of torch.no_grad().
Best.
K. Frank