diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e355010..eb5807d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -36,3 +36,8 @@ jobs:
     - name: Test with pytest
       run: |
         docker run --rm cvdm-package:${{ matrix.python-version }} pytest
+
+    # Step 5: Run the container and execute tests using mypy
+    - name: Test with mypy
+      run: |
+        docker run --rm cvdm-package:${{ matrix.python-version }} mypy .
diff --git a/cvdm/utils/inference_utils.py b/cvdm/utils/inference_utils.py
index 0d774b1..0c5c396 100644
--- a/cvdm/utils/inference_utils.py
+++ b/cvdm/utils/inference_utils.py
@@ -90,6 +90,7 @@ def create_output_montage(
             (pred_y, y, gamma_vec),
             axis=2,
         )
+    print(concatenated_images.shape)
     image: np.ndarray = montage(
         np.squeeze(concatenated_images),
         channel_axis=channel_axis,