You error message seems to raise a NaN loss or do you also see the “not implemented for Byte” error?
In the former case, could you create a new topic with a description of your model and use case (e.g. are you using mixed-precision training, what type or model etc.) and tag me there please?



 mapostig:

Moreover the image visualization after the tensor cast is very poor.


If you are visualizing with matplotlib, you might have to check the current range, as matplotlib might try to normalize and clip the image.