compute ms ssim in test steps

broadinstitute · Jan 21, 2025 · 591ec1b · 591ec1b
1 parent dec6b04
commit 591ec1b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
@@ -887,7 +887,7 @@ def test_step(self, batch):
             )
             self.kid.update_state(images, generated_images)
             max_pixel_value = tf.reduce_max(tf.abs(images))
-            max_val = 2 * max_pixel_value.numpy()  # Double the max absolute value
+            max_val = 2 * max_pixel_value  # Double the max absolute value
             self.ms_ssim.update_state(images, generated_images, max_val)
 
         return {m.name: m.result() for m in self.metrics}