From ce5f479f9f2691f1708ecec9b714167fa0e23272 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:09:51 -0800 Subject: [PATCH] try casting --- core/leras/layers/MsSsim.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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