diff --git a/cosmos_predict2/robot_multiview.py b/cosmos_predict2/robot_multiview.py
index e9f6ddfe0..ff3f03166 100644
--- a/cosmos_predict2/robot_multiview.py
+++ b/cosmos_predict2/robot_multiview.py
@@ -164,7 +164,8 @@ def __init__(
                 # pyrefly: ignore  # implicit-import
                 transforms.v2.Resize(size=(self.height, self.width), antialias=True),
                 # pyrefly: ignore  # implicit-import
-                transforms.v2.ToTensor(),
+                transforms.v2.ToImage(),
+                transforms.v2.ToDtype(torch.float32, scale=True),
                 # pyrefly: ignore  # implicit-import
                 transforms.v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
             ]