diff --git a/bob/ip/binseg/engine/predictor.py b/bob/ip/binseg/engine/predictor.py
index 33a91d0b7bc97daf583a42cdbe7ad7328bfafa88..804352eb86617128c92c1ec6fcd15c0c448d3d9d 100644
--- a/bob/ip/binseg/engine/predictor.py
+++ b/bob/ip/binseg/engine/predictor.py
@@ -100,7 +100,7 @@ def _save_overlayed_png(stem, image, prob, output_folder):
     """
 
     image = VF.to_pil_image(image)
-    prob = VF.to_pil_image(prob.cpu().squeeze(0))
+    prob = VF.to_pil_image(prob.cpu())
 
     # color and overlay
     prob_green = PIL.ImageOps.colorize(prob, (0, 0, 0), (0, 255, 0))