diff --git a/lensless/utils/dataset.py b/lensless/utils/dataset.py
index 3ea548dd..14dd2629 100644
--- a/lensless/utils/dataset.py
+++ b/lensless/utils/dataset.py
@@ -1050,7 +1050,7 @@ def _get_images_pair(self, idx):
         if self.alignment is not None:
             lensed_np = resize(
                 lensed_np,
-                shape=(self.alignment["height"], self.alignment["width"]),
+                shape=(self.alignment["height"], self.alignment["width"], 3),
                 interpolation=cv2.INTER_NEAREST,
             )
         elif self.display_res is not None: