compute ms ssim in test steps

broadinstitute · Jan 21, 2025 · e12001a · e12001a
1 parent 70363ad
commit e12001a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
@@ -886,7 +886,7 @@ def test_step(self, batch):
                 num_images=self.batch_size, diffusion_steps=20
             )
             self.kid.update_state(images, generated_images)
-            max_pixel_value = tf.reduce_max(tf.abs(images))
+            max_pixel_value = tf.reduce_max(tf.abs(generated_images))
             max_val = 2 * max_pixel_value  # Double the max absolute value
             self.ms_ssim.update_state(images, generated_images, max_val)