diff --git a/configs/data/im2im/segmentation_plugin.yaml b/configs/data/im2im/segmentation_plugin.yaml
index f296e469..1ffa15e1 100644
--- a/configs/data/im2im/segmentation_plugin.yaml
+++ b/configs/data/im2im/segmentation_plugin.yaml
@@ -151,6 +151,7 @@ transforms:
   predict:
     _target_: monai.transforms.Compose
     transforms:
+      - _target_: cyto_dl.datamodules.dataframe.utils.RemoveNaNKeysd
       # load
       - _target_: monai.transforms.LoadImaged
         keys: ${source_col}