diff --git a/src/mednet/libs/segmentation/models/typing.py b/src/mednet/libs/segmentation/models/typing.py
index 5a24afb6f4857480fdad74d92f8ab3eb199581d0..11ec1457ff1ece9fa26d779190c529c0bd78c724 100644
--- a/src/mednet/libs/segmentation/models/typing.py
+++ b/src/mednet/libs/segmentation/models/typing.py
@@ -9,5 +9,7 @@ import typing
 Checkpoint: typing.TypeAlias = typing.MutableMapping[str, typing.Any]
 """Definition of a lightning checkpoint."""
 
-SegmentationPrediction: typing.TypeAlias = tuple[pathlib.Path, pathlib.Path]
-"""The sample name, the target, and the predicted value."""
+SegmentationPrediction: typing.TypeAlias = tuple[
+    pathlib.Path, pathlib.Path, pathlib.Path, pathlib.Path
+]
+"""The sample name, the target, mask, and the prediction."""