Update main.py

main.py

@@ -117,6 +117,8 @@ def predict_next_frame(previous_frames: List[np.ndarray], previous_actions: List
     
     # Convert the generated frame to the correct format
     #new_frame = (new_frame * 255).astype(np.uint8).transpose(1, 2, 0)
+    new_frame = new_frame.transpose(1, 2, 0)
+
     
     # Resize the frame to 256x256 if necessary
     #if new_frame.shape[:2] != (height, width):