SSIM metric not working with half-precision

[ncuxomun]

Hi!

I am using both MSE and SSIM metrics in one (total) loss function in an AutoEncoder but I am not able to use the half-precision on my GPU because of the error as shown below:

What I understand is that my input X (from the batch) is on Float32 but my pred_X is on Float16 (I think it is due to the half-precision), hence these two are incompatible resulting in error. Interestingly, they are OK when computed in MSE loss, but NOT OK when computed in SSIM loss.

How can I resolve the issue if I want to use half-precision on my GPU?
Thank you for your attention!