[bugfix] Double Normalization in Preprocessing Dataset (#1055)

2026-01-31 11:11:59 -08:00
parent 38a6bd93d3
commit caa1c402ba
1 changed files with 3 additions and 1 deletions
@@ -366,9 +366,11 @@ class ImageTransformStage(DatasetStage):
            image = self.transform_topcrop(image)
        elif self.transform is not None:
            image = self.transform(image)
+            image = image.float() / 127.5 - 1.0
+        else:
+            image = image.float() / 127.5 - 1.0

        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-        image = image.float() / 127.5 - 1.0
        batch.pixel_values = image
        return batch