Update utils/data.py

utils/data.py

@@ -177,7 +177,7 @@ def preprocess(
         inputs[0, :, :, :] = inputs[0, :, :, :] - inputs[0, :, 0, 0].mean(axis=0)
         inputs[1, :, :, :] = inputs[1, :, :, :] - inputs[1, :, 0, 0].mean(axis=0)
 
-    return inputs[np.newaxis, :].astype(np.float32)
+    return np.squeeze(inputs).transpose(1, 2, 0).astype(np.float32)
 
 
 def random_sample_np(data: np.ndarray, size: int) -> np.ndarray: