diff --git a/tests/test_padchest.py b/tests/test_padchest.py
index 5fc342e1ab8d61ec55adaf41c7fe279d19499431..98eaeb2d35036727eab39cd85578d652966c1b36 100644
--- a/tests/test_padchest.py
+++ b/tests/test_padchest.py
@@ -60,18 +60,7 @@ def test_loading(database_checkers, name: str, dataset: str, num_labels: int):
     datamodule.setup("predict")  # sets up all datasets
 
     if dataset in datamodule.predict_dataloader():
-        # Not all datamodules have a test set
-        try:
-            loader = datamodule.predict_dataloader()[dataset]
-        except KeyError as e:
-            if str(e) == "test" and name in [
-                "idiap",
-                "no_tb_idiap",
-                "cardiomegaly_idiap",
-            ]:
-                assert True
-            else:
-                assert False
+        loader = datamodule.predict_dataloader()[dataset]
 
         limit = 3  # limit load checking
         for batch in loader: