diff --git a/Matrix-Game-2/inference.py b/Matrix-Game-2/inference.py
index f4df03b..3201b6d 100644
--- a/Matrix-Game-2/inference.py
+++ b/Matrix-Game-2/inference.py
@@ -40,7 +40,8 @@ def __init__(self, args):
 
         self.frame_process = v2.Compose([
             v2.Resize(size=(352, 640), antialias=True),
-            v2.ToTensor(),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
             v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
         ])