diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py
index d9533f4..6fa7c69 100644
--- a/core/leras/layers/MsSsim.py
+++ b/core/leras/layers/MsSsim.py
@@ -17,10 +17,9 @@ class MsSsim(nn.LayerBase):
 
     def __call__(self, y_true, y_pred, max_val):
         # Transpose images from NCHW to NHWC
-        y_true_t = tf.transpose(y_true, [0, 2, 3, 1])
-        y_pred_t = tf.transpose(y_pred, [0, 2, 3, 1])
-        with tf.device('/CPU:0'):
-            loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors)
+        y_true_t = tf.transpose(tf.cast(y_true), [0, 2, 3, 1])
+        y_pred_t = tf.transpose(tf.cast(y_pred), [0, 2, 3, 1])
+        loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors)
         return (1.0 - loss) / 2.0
 
 nn.MsSsim = MsSsim