diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
index f358da710..486287ed5 100644
--- a/ml4h/models/diffusion_blocks.py
+++ b/ml4h/models/diffusion_blocks.py
@@ -888,7 +888,7 @@ def test_step(self, batch):
             self.kid.update_state(images, generated_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)
+            self.ms_ssim.update_state(images, generated_images, 255)
 
         return {m.name: m.result() for m in self.metrics}