diff --git a/src/mednet/libs/classification/config/data/montgomery/datamodule.py b/src/mednet/libs/classification/config/data/montgomery/datamodule.py
index 4a61df7df30d507b4be259ce2e57d7ca91c37b03..52f2332053ffa105b36f353f95190b84d689a21b 100644
--- a/src/mednet/libs/classification/config/data/montgomery/datamodule.py
+++ b/src/mednet/libs/classification/config/data/montgomery/datamodule.py
@@ -13,10 +13,10 @@ import PIL.Image
 from mednet.libs.classification.data.typing import (
     ClassificationRawDataLoader as _ClassificationRawDataLoader,
 )
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.image_utils import remove_black_borders
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
 from torchvision.transforms.functional import to_tensor
 
diff --git a/src/mednet/libs/classification/config/data/nih_cxr14/datamodule.py b/src/mednet/libs/classification/config/data/nih_cxr14/datamodule.py
index 7712b4fe8aa864e6d27e033157552719b09bcdef..719a78adff1e4517af05a4cbd962791c1aa6dc06 100644
--- a/src/mednet/libs/classification/config/data/nih_cxr14/datamodule.py
+++ b/src/mednet/libs/classification/config/data/nih_cxr14/datamodule.py
@@ -13,9 +13,9 @@ import PIL.Image
 from mednet.libs.classification.data.typing import (
     ClassificationRawDataLoader as _ClassificationRawDataLoader,
 )
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
 from torchvision.transforms.functional import to_tensor
 
diff --git a/src/mednet/libs/classification/config/data/padchest/datamodule.py b/src/mednet/libs/classification/config/data/padchest/datamodule.py
index cc9bd7349a449f55cf9e134c4e8602b3645a90bb..20f6e221ae12f7778d125788d101161ac7c5e494 100644
--- a/src/mednet/libs/classification/config/data/padchest/datamodule.py
+++ b/src/mednet/libs/classification/config/data/padchest/datamodule.py
@@ -14,10 +14,10 @@ import PIL.Image
 from mednet.libs.classification.data.typing import (
     ClassificationRawDataLoader as _ClassificationRawDataLoader,
 )
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.image_utils import remove_black_borders
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
 from torchvision.transforms.functional import to_tensor
 
diff --git a/src/mednet/libs/classification/config/data/shenzhen/datamodule.py b/src/mednet/libs/classification/config/data/shenzhen/datamodule.py
index ba2ccc5ce2956fe376db7578cb413c0e52320fd4..255d70b0eec3d8ca6ec51bf28eef983a0828db2f 100644
--- a/src/mednet/libs/classification/config/data/shenzhen/datamodule.py
+++ b/src/mednet/libs/classification/config/data/shenzhen/datamodule.py
@@ -13,10 +13,10 @@ import PIL.Image
 from mednet.libs.classification.data.typing import (
     ClassificationRawDataLoader as _ClassificationRawDataLoader,
 )
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.image_utils import remove_black_borders
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
 from torchvision.transforms.functional import to_tensor
 
diff --git a/src/mednet/libs/classification/config/data/tbpoc/datamodule.py b/src/mednet/libs/classification/config/data/tbpoc/datamodule.py
index 8f46b3beae6e6533bdb8354286737b7a87cf73fc..7d9c9bd8913fe5fb3eaf3ad53c1a309db3dd7122 100644
--- a/src/mednet/libs/classification/config/data/tbpoc/datamodule.py
+++ b/src/mednet/libs/classification/config/data/tbpoc/datamodule.py
@@ -9,10 +9,10 @@ import PIL.Image
 from mednet.libs.classification.data.typing import (
     ClassificationRawDataLoader as _ClassificationRawDataLoader,
 )
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.image_utils import remove_black_borders
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
 from torchvision.transforms.functional import to_tensor
 
diff --git a/src/mednet/libs/classification/config/data/tbx11k/datamodule.py b/src/mednet/libs/classification/config/data/tbx11k/datamodule.py
index 9cef49b1ea9960821b12674fad41e6fa650ffae2..f120a2afc4216312c1bf0f88931273c6b2045622 100644
--- a/src/mednet/libs/classification/config/data/tbx11k/datamodule.py
+++ b/src/mednet/libs/classification/config/data/tbx11k/datamodule.py
@@ -13,9 +13,9 @@ import typing_extensions
 from mednet.libs.classification.data.typing import (
     ClassificationRawDataLoader as _ClassificationRawDataLoader,
 )
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
 from torch.utils.data._utils.collate import default_collate_fn_map
 from torchvision.transforms.functional import to_tensor
diff --git a/src/mednet/libs/classification/data/typing.py b/src/mednet/libs/classification/data/typing.py
index 0d18220e0f54b339cf296278eba7c9c38df6fd86..10dc375be754fd0a7a0cb541d678ab3d341fd6de 100644
--- a/src/mednet/libs/classification/data/typing.py
+++ b/src/mednet/libs/classification/data/typing.py
@@ -1,6 +1,9 @@
 import typing
 
-from mednet.libs.common.data.typing import RawDataLoader, Sample
+import torch
+from mednet.libs.common.data.typing import RawDataLoader
+
+Sample: typing.TypeAlias = tuple[torch.Tensor, typing.Mapping[str, typing.Any]]
 
 
 class ClassificationRawDataLoader(RawDataLoader):
diff --git a/src/mednet/libs/classification/engine/saliency/completeness.py b/src/mednet/libs/classification/engine/saliency/completeness.py
index 2eac54295f0bdbdb1151dd290db1bbce2ddb9661..f632afee5991ec1ea3e8b5a1a4612fda33a024f8 100644
--- a/src/mednet/libs/classification/engine/saliency/completeness.py
+++ b/src/mednet/libs/classification/engine/saliency/completeness.py
@@ -11,7 +11,7 @@ import lightning.pytorch
 import numpy as np
 import torch
 import tqdm
-from mednet.libs.common.data.typing import Sample
+from mednet.libs.classification.data.typing import Sample
 from mednet.libs.common.engine.device import DeviceManager
 from pytorch_grad_cam.metrics.road import (
     ROADLeastRelevantFirstAverage,
diff --git a/src/mednet/libs/common/data/typing.py b/src/mednet/libs/common/data/typing.py
index 521c8026f57dcc797bae31c239a3357a123f835c..455cb1b37647ea55b02f5334131509e3eca7b85b 100644
--- a/src/mednet/libs/common/data/typing.py
+++ b/src/mednet/libs/common/data/typing.py
@@ -9,7 +9,9 @@ import typing
 import torch
 import torch.utils.data
 
-Sample: typing.TypeAlias = tuple[torch.Tensor, typing.Mapping[str, typing.Any]]
+Sample: typing.TypeAlias = tuple[
+    torch.Tensor | typing.Mapping[str, torch.Tensor], typing.Mapping[str, typing.Any]
+]
 """Definition of a sample.
 
 First parameter
diff --git a/src/mednet/libs/segmentation/config/data/avdrive/datamodule.py b/src/mednet/libs/segmentation/config/data/avdrive/datamodule.py
index 83156a2c03dbf2fadf9054b9d32d6c31fe0502f8..5567d27cc5ad62ea1dc4631f0759ab2d2f42600a 100644
--- a/src/mednet/libs/segmentation/config/data/avdrive/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/avdrive/datamodule.py
@@ -9,9 +9,9 @@ import pathlib
 import PIL.Image
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/chasedb1/datamodule.py b/src/mednet/libs/segmentation/config/data/chasedb1/datamodule.py
index 0cfab7ee7b66ed27f607559f4ff3bf199550b9ca..16a12166647238ae09ff0a9085592c832d3eda05 100644
--- a/src/mednet/libs/segmentation/config/data/chasedb1/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/chasedb1/datamodule.py
@@ -10,9 +10,9 @@ import PIL.Image
 import pkg_resources
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/cxr8/datamodule.py b/src/mednet/libs/segmentation/config/data/cxr8/datamodule.py
index acc6c83bf0b3112a01d12cdebd8a615b4eabd5c1..fee56bf01925c2f6d2a5891449b0fe256b94dcd6 100644
--- a/src/mednet/libs/segmentation/config/data/cxr8/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/cxr8/datamodule.py
@@ -11,8 +11,8 @@ import PIL.Image
 import torch
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/drhagis/datamodule.py b/src/mednet/libs/segmentation/config/data/drhagis/datamodule.py
index 153c0abe36023d4e57c147dbb1c43f35b42c816e..1b0807dd4d3e94a2412760bdcaeb6c1a080c3570 100644
--- a/src/mednet/libs/segmentation/config/data/drhagis/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/drhagis/datamodule.py
@@ -9,9 +9,9 @@ import pathlib
 import PIL.Image
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/drionsdb/datamodule.py b/src/mednet/libs/segmentation/config/data/drionsdb/datamodule.py
index 41cc7fb8f02309bcb6d982e70e0d403bf5529ec6..5e7076afe01c983c29189f10dbe81874e4d1bc4b 100644
--- a/src/mednet/libs/segmentation/config/data/drionsdb/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/drionsdb/datamodule.py
@@ -11,9 +11,9 @@ import PIL.Image
 import pkg_resources
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/drishtigs1/datamodule.py b/src/mednet/libs/segmentation/config/data/drishtigs1/datamodule.py
index 1c65e40dce56d6066105808c61a037211770a13f..2ddc1b83ca4e867b9ee1ad9bcc17a154040b815a 100644
--- a/src/mednet/libs/segmentation/config/data/drishtigs1/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/drishtigs1/datamodule.py
@@ -10,9 +10,9 @@ import PIL.Image
 import pkg_resources
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/drive/datamodule.py b/src/mednet/libs/segmentation/config/data/drive/datamodule.py
index 176293431877ed2dcf6e124c90faf4de0f7944a6..01639da4c077e3dfc3251b084184920c6aa509e1 100644
--- a/src/mednet/libs/segmentation/config/data/drive/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/drive/datamodule.py
@@ -9,9 +9,9 @@ import pathlib
 import PIL.Image
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/hrf/datamodule.py b/src/mednet/libs/segmentation/config/data/hrf/datamodule.py
index 811ac24f9bc7244ffda037ef1c93b3fda8e8df2a..ed7ac302732c249c2a11dd404ff41e3392e07f3a 100644
--- a/src/mednet/libs/segmentation/config/data/hrf/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/hrf/datamodule.py
@@ -9,9 +9,9 @@ import pathlib
 import PIL.Image
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/iostar/datamodule.py b/src/mednet/libs/segmentation/config/data/iostar/datamodule.py
index ff67de9e272c44ec2c594a9edb5a11382417af84..60283259d5a9b73001f1e66c157d8c4cc856877f 100644
--- a/src/mednet/libs/segmentation/config/data/iostar/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/iostar/datamodule.py
@@ -9,9 +9,9 @@ import pathlib
 import PIL.Image
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/jsrt/datamodule.py b/src/mednet/libs/segmentation/config/data/jsrt/datamodule.py
index 60f185b52a42974ddc29b8a247fe4139bc9b6162..ada65db973f9e75262e4dc2b7b9fa39c7369878c 100644
--- a/src/mednet/libs/segmentation/config/data/jsrt/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/jsrt/datamodule.py
@@ -12,8 +12,8 @@ import skimage.exposure
 import torch
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/montgomery/datamodule.py b/src/mednet/libs/segmentation/config/data/montgomery/datamodule.py
index 5792d941328fae7f18238bc73a4133b4cd2c0e1e..663f9371a8191e6db6051aa389f0e919cf55c6b1 100644
--- a/src/mednet/libs/segmentation/config/data/montgomery/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/montgomery/datamodule.py
@@ -12,8 +12,8 @@ import pkg_resources
 import torch
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/refuge/datamodule.py b/src/mednet/libs/segmentation/config/data/refuge/datamodule.py
index a7bcd24d161df05c2c104571340b30271793a7c7..696a204e2fa625950d8d58633c04acff9dcb7391 100644
--- a/src/mednet/libs/segmentation/config/data/refuge/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/refuge/datamodule.py
@@ -10,9 +10,9 @@ import PIL.Image
 import pkg_resources
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/rimoner3/datamodule.py b/src/mednet/libs/segmentation/config/data/rimoner3/datamodule.py
index e4e9e5815c83c93afe95ed331c2489fa90544101..815744418bf905b825a063c03ed687f78acc0385 100644
--- a/src/mednet/libs/segmentation/config/data/rimoner3/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/rimoner3/datamodule.py
@@ -10,8 +10,8 @@ import PIL.Image
 import pkg_resources
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/shenzhen/datamodule.py b/src/mednet/libs/segmentation/config/data/shenzhen/datamodule.py
index f0be322abb3a91607580c9bd686f844d9d824984..a65f3239dd752ae985a785db90345f716cd0ded4 100644
--- a/src/mednet/libs/segmentation/config/data/shenzhen/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/shenzhen/datamodule.py
@@ -10,9 +10,9 @@ import PIL.Image
 import torch
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/config/data/stare/datamodule.py b/src/mednet/libs/segmentation/config/data/stare/datamodule.py
index 9c8247823eb51532a8ef242e97a9d5fc4c12b96f..06bceb509f6be8c8f0a8dd62d1bee1603b6e7d98 100644
--- a/src/mednet/libs/segmentation/config/data/stare/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/stare/datamodule.py
@@ -10,9 +10,9 @@ import PIL.Image
 import pkg_resources
 from mednet.libs.common.data.datamodule import CachingDataModule
 from mednet.libs.common.data.split import make_split
-from mednet.libs.common.data.typing import Sample
 from mednet.libs.common.models.transforms import crop_image_to_mask
 from mednet.libs.common.utils.rc import load_rc
+from mednet.libs.segmentation.data.typing import Sample
 from mednet.libs.segmentation.data.typing import (
     SegmentationRawDataLoader as _SegmentationRawDataLoader,
 )
diff --git a/src/mednet/libs/segmentation/data/typing.py b/src/mednet/libs/segmentation/data/typing.py
index 626464fdbe222e8008c4cb213b82fc0aa54a7e3d..af27de404771971c45e3f684783a8f1b0ff7a807 100644
--- a/src/mednet/libs/segmentation/data/typing.py
+++ b/src/mednet/libs/segmentation/data/typing.py
@@ -1,4 +1,11 @@
-from mednet.libs.common.data.typing import RawDataLoader, Sample
+import typing
+
+import torch
+from mednet.libs.common.data.typing import RawDataLoader
+
+Sample: typing.TypeAlias = tuple[
+    typing.Mapping[str, torch.Tensor], typing.Mapping[str, typing.Any]
+]
 
 
 class SegmentationRawDataLoader(RawDataLoader):