diff --git a/src/mednet/libs/segmentation/config/data/drive/datamodule.py b/src/mednet/libs/segmentation/config/data/drive/datamodule.py
index 893cbbb8682fa248cc5eaf642ed404ad25d804d6..d7f2890ac96b8c8b4e80429410433927e706e6e0 100644
--- a/src/mednet/libs/segmentation/config/data/drive/datamodule.py
+++ b/src/mednet/libs/segmentation/config/data/drive/datamodule.py
@@ -25,7 +25,7 @@ database."""
 
 
 class SegmentationRawDataLoader(_SegmentationRawDataLoader):
-    """A specialized raw-data-loader for the Montgomery dataset."""
+    """A specialized raw-data-loader for the Drive dataset."""
 
     datadir: str
     """This variable contains the base directory where the database raw data is
@@ -73,7 +73,7 @@ class SegmentationRawDataLoader(_SegmentationRawDataLoader):
 
 
 def make_split(basename: str) -> DatabaseSplit:
-    """Return a database split for the Montgomery database.
+    """Return a database split for the Drive database.
 
     Parameters
     ----------
diff --git a/src/mednet/libs/segmentation/tests/test_drive.py b/src/mednet/libs/segmentation/tests/test_drive.py
index a90f63184cf1891b1d3f27a948d1ae8c6a84230e..d1a1301d47522459397e5b7b1b22b733a42faeb2 100644
--- a/src/mednet/libs/segmentation/tests/test_drive.py
+++ b/src/mednet/libs/segmentation/tests/test_drive.py
@@ -109,17 +109,10 @@ def test_raw_transforms_image_quality(database_checkers, datadir):
     ["lwnet"],
 )
 def test_model_transforms_image_quality(database_checkers, datadir, model_name):
-    # Densenet's model.name is "densenet-212" and does not correspond to its module name.
-    if model_name == "densenet":
-        reference_histogram_file = str(
-            datadir
-            / "histograms/models/histograms_densenet-121_drive_default.json",
-        )
-    else:
-        reference_histogram_file = str(
-            datadir
-            / f"histograms/models/histograms_{model_name}_drive_default.json",
-        )
+    reference_histogram_file = str(
+        datadir
+        / f"histograms/models/histograms_{model_name}_drive_default.json",
+    )
 
     datamodule = importlib.import_module(
         ".default",
diff --git a/src/mednet/libs/segmentation/tests/test_stare.py b/src/mednet/libs/segmentation/tests/test_stare.py
index 99b36913fae67d5d75178d2383e6430898d0509a..9b0038945da8c6648664dcd124add57eed898cf5 100644
--- a/src/mednet/libs/segmentation/tests/test_stare.py
+++ b/src/mednet/libs/segmentation/tests/test_stare.py
@@ -111,16 +111,9 @@ def test_raw_transforms_image_quality(database_checkers, datadir):
     ["lwnet"],
 )
 def test_model_transforms_image_quality(database_checkers, datadir, model_name):
-    # Densenet's model.name is "densenet-212" and does not correspond to its module name.
-    if model_name == "densenet":
-        reference_histogram_file = str(
-            datadir / "histograms/models/histograms_densenet-121_stare_ah.json",
-        )
-    else:
-        reference_histogram_file = str(
-            datadir
-            / f"histograms/models/histograms_{model_name}_stare_ah.json",
-        )
+    reference_histogram_file = str(
+        datadir / f"histograms/models/histograms_{model_name}_stare_ah.json",
+    )
 
     datamodule = importlib.import_module(
         ".ah",