Do you only get the NaN values using one of my implementations or also your original one using sklearn?
If just using mine, could you post an input tensor, which creates these NaN values?
It’s a bit weird, as I’ve tested your approach to mine and got the same results (up to floating point precision).
So if the scaling creates NaNs in your training, both approaches should behave the same.



 lorenzo_fabbri:

I finally realized that if I do not use your scaler, I get the right losses!


Are you getting the right losses using the sklearn approach now?



 lorenzo_fabbri:

Moreover, it’s much faster if I do not scale the images (from roughly 19 minutes to 12).


This is expected, as you will save some operations or am I missing something?