Manage image resizing

handler.py

@@ -64,9 +64,16 @@ class EndpointHandler():
         
         # 4. Process and select the best mask
         try:
+            pred_masks_raw = outputs.pred_masks.cpu()
+
+            # The model may output 5-dim tensors, but the post-processor expects 4-dim.
+            # We squeeze the extra dimension to fix this.
+            if pred_masks_raw.ndim == 5:
+                pred_masks_raw = pred_masks_raw.squeeze(1)
+
             # Use the processor's post-processing utility to resize masks and remove padding
             masks = self.processor.post_process_masks(
-                outputs.pred_masks.cpu(), 
+                pred_masks_raw, 
                 inputs["original_sizes"].cpu(), 
                 inputs["reshaped_input_sizes"].cpu()
             )[0]