diff --git a/tests/test_cli.py b/tests/test_cli.py
index 4ee0e2c6ab454161dbf871030b30b989593d1c85..a512d6add392dc09439daaca7e0d8f9b01fbc443 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -241,8 +241,7 @@ def test_train_pasa_montgomery(temporary_basedir):
         keywords = {
             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"^Balancing samples from dataset using metadata targets `label`$": 1,
+            r"^Applying train/valid loss balancing...$": 1,
             r"^Training for at most 1 epochs.$": 1,
             r"^Uninitialised pasa model - computing z-norm factors from train dataloader.$": 1,
             r"^Writing run metadata at.*$": 1,
@@ -323,8 +322,7 @@ def test_train_pasa_montgomery_from_checkpoint(temporary_basedir):
         keywords = {
             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"^Balancing samples from dataset using metadata targets `label`$": 1,
+            r"^Applying train/valid loss balancing...$": 1,
             r"^Training for at most 2 epochs.$": 1,
             r"^Resuming from epoch 0 \(checkpoint file: .*$": 1,
             r"^Writing run metadata at.*$": 1,