diff --git a/configs/experiment/im2im/segmentation_plugin.yaml b/configs/experiment/im2im/segmentation_plugin.yaml
index deb8513d..e89c9b58 100644
--- a/configs/experiment/im2im/segmentation_plugin.yaml
+++ b/configs/experiment/im2im/segmentation_plugin.yaml
@@ -78,5 +78,5 @@ callbacks:
     _target_: cyto_dl.callbacks.ImageSaver
     save_dir: ${paths.output_dir}
     save_every_n_epochs: ${model.save_images_every_n_epochs}
-    stages: ["train", "test", "val"]
+    stages: ["train", "test", "val", "predict"]
     save_input: True