diff --git a/tests/test_cli.py b/tests/test_cli.py
index 64a9e2d5f8bdd9fed55e156fe185f8d9ca20c0ec..4828e1d09ca51d1b06c5cb3c29fb1f8b62ddccd0 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -224,7 +224,7 @@ def test_train_pasa_montgomery(temporary_basedir):
             r"^Writing command-line for reproduction at .*$": 1,
             r"^Loading dataset:`train` without caching. Trade-off: CPU RAM usage: less | Disk I/O: more.$": 1,
             r"^Loading dataset:`validation` without caching. Trade-off: CPU RAM usage: less | Disk I/O: more.$": 1,
-            r"^Applying datamodule train sampler balancing...$": 1,
+            r"^Applying DataModule train sampler balancing...$": 1,
             r"^Balancing samples from dataset using metadata targets `label`$": 1,
             r"^Training for at most 1 epochs.$": 1,
             r"^Uninitialised pasa model - computing z-norm factors from train dataloader.$": 1,
@@ -321,7 +321,7 @@ def test_train_pasa_montgomery_from_checkpoint(temporary_basedir):
             r"^Writing command-line for reproduction at .*$": 1,
             r"^Loading dataset:`train` without caching. Trade-off: CPU RAM usage: less | Disk I/O: more.$": 1,
             r"^Loading dataset:`validation` without caching. Trade-off: CPU RAM usage: less | Disk I/O: more.$": 1,
-            r"^Applying datamodule train sampler balancing...$": 1,
+            r"^Applying DataModule train sampler balancing...$": 1,
             r"^Balancing samples from dataset using metadata targets `label`$": 1,
             r"^Training for at most 2 epochs.$": 1,
             r"^Resuming from epoch 0 \(checkpoint file: .*$": 1,