diff --git a/tests/test_cli.py b/tests/test_cli.py
index 796a795ccfc5ff09b61fff9c563655ca0e2795a0..1204257d4e4e06b08914f02f655b880ac1ea8034 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -328,7 +328,7 @@ def test_predict_pasa_montgomery(temporary_basedir, datadir):
                 "-vv",
                 "--batch-size=1",
                 "--relevance-analysis",
-                f"--weight={str(datadir / 'lfs' / 'models' / 'pasa.pth')}",
+                f"--weight={str(datadir / 'lfs' / 'models' / 'pasa.ckpt')}",
                 f"--output-folder={output_folder}",
             ],
         )
@@ -346,7 +346,6 @@ def test_predict_pasa_montgomery(temporary_basedir, datadir):
 
         keywords = {
             r"^Loading checkpoint from.*$": 1,
-            r"^Total time:.*$": 3,
             r"^Relevance analysis.*$": 3,
         }
         buf.seek(0)
@@ -564,7 +563,7 @@ def test_predict_signstotb_montgomery_rs(temporary_basedir, datadir):
                 "-vv",
                 "--batch-size=1",
                 "--relevance-analysis",
-                f"--weight={str(datadir / 'lfs' / 'models' / 'signstotb.pth')}",
+                f"--weight={str(datadir / 'lfs' / 'models' / 'signstotb.ckpt')}",
                 f"--output-folder={output_folder}",
             ],
         )
@@ -582,7 +581,6 @@ def test_predict_signstotb_montgomery_rs(temporary_basedir, datadir):
 
         keywords = {
             r"^Loading checkpoint from.*$": 1,
-            r"^Total time:.*$": 3 * 15,
             r"^Starting relevance analysis for subset.*$": 3,
             r"^Creating and saving plot at.*$": 3,
         }
@@ -665,7 +663,7 @@ def test_predict_logreg_montgomery_rs(temporary_basedir, datadir):
                 "montgomery_rs",
                 "-vv",
                 "--batch-size=1",
-                f"--weight={str(datadir / 'lfs' / 'models' / 'logreg.pth')}",
+                f"--weight={str(datadir / 'lfs' / 'models' / 'logreg.ckpt')}",
                 f"--output-folder={output_folder}",
             ],
         )
@@ -679,7 +677,6 @@ def test_predict_logreg_montgomery_rs(temporary_basedir, datadir):
 
         keywords = {
             r"^Loading checkpoint from.*$": 1,
-            r"^Total time:.*$": 3,
             r"^Logistic regression identified: saving model weights.*$": 1,
         }
         buf.seek(0)