Fix SSIM eval for bf16

Summary: As titled, similar to previous diff, fix SSIM eval for data type as well Reviewed By: dongyuzheng Differential Revision: D51223782 fbshipit-source-id: 74e6d2b8282e43f2dd8a0a0b87b4bd1f8405ae33
pytorch · Nov 13, 2023 · f6325a2 · f6325a2
1 parent 5faeeee
commit f6325a2
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torcheval/metrics/image/ssim.py b/torcheval/metrics/image/ssim.py
@@ -55,6 +55,9 @@ def update(
         """
         if images_1.shape != images_2.shape:
             raise RuntimeError("The two sets of images must have the same shape.")
+        # convert to fp32, mostly for bf16 types
+        images_1 = images_1.to(dtype=torch.float32)
+        images_2 = images_2.to(dtype=torch.float32)
 
         batch_size = images_1.shape[0]