adding an assert

ashesh-0 · Jan 15, 2025 · 63830fb · 63830fb
1 parent 38024ee
commit 63830fb
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/predtiler/tile_stitcher.py b/src/predtiler/tile_stitcher.py
@@ -15,8 +15,7 @@ def stitch_predictions(predictions:np.ndarray, manager):
     mng = manager
     shape = list(mng.data_shape)
     shape.append(predictions.shape[1])
-    print(shape)
-
+    assert mng.patch_shape[-2:] == predictions.shape[-2:], 'Patch shape and predictions shape must match. Please set the patch shape correctly'
     output = np.zeros(shape, dtype=predictions.dtype)
     for dset_idx in range(predictions.shape[0]):
         # grid start, grid end