From 0225a29832c1a35eacf080e3ea5875a3473bea8a Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 18:07:27 -0800 Subject: [PATCH] cpu --- core/leras/layers/MsSsim.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/leras/layers/MsSsim.py b/core/leras/layers/MsSsim.py index 12b9ea1..d9533f4 100644 --- a/core/leras/layers/MsSsim.py +++ b/core/leras/layers/MsSsim.py @@ -19,7 +19,8 @@ class MsSsim(nn.LayerBase): # 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]) - loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) + with tf.device('/CPU:0'): + 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