diff --git a/segmentation/cells/UNet3DPlatyCellProbs.model/UNet3DPlatyCellProbs.model.yaml b/segmentation/cells/UNet3DPlatyCellProbs.model/UNet3DPlatyCellProbs.model.yaml
index 92ec3644..cf062dd9 100644
--- a/segmentation/cells/UNet3DPlatyCellProbs.model/UNet3DPlatyCellProbs.model.yaml
+++ b/segmentation/cells/UNet3DPlatyCellProbs.model/UNet3DPlatyCellProbs.model.yaml
@@ -39,15 +39,10 @@ outputs:
         reference_input: raw
         scale: [1, 1, 1, 1, 1]
         offset: [0, 0, 0, 0, 0]
-    # TODO for now we could put this into the model
-    # TODO we need a post-processing op to average the output channels
-    # postprocessing:
-    #     - name: average_channels
-    #       kwargs: {start_channel: 0, stop_channel: 3}
 
 language: python
 framework: pytorch
-source: models.UNetAnisotropic
+source: ./models.py::UNetAnisotropic
 kwargs:
   in_channels: 1
   out_channels: 12
diff --git a/segmentation/cells/UNet3DPlatyCellProbs.model/test_output.npy b/segmentation/cells/UNet3DPlatyCellProbs.model/test_output.npy
index 4b51ade2..01c96313 100644
Binary files a/segmentation/cells/UNet3DPlatyCellProbs.model/test_output.npy and b/segmentation/cells/UNet3DPlatyCellProbs.model/test_output.npy differ