diff --git a/cellseg_models_pytorch/training/tests/test_training.py b/cellseg_models_pytorch/training/tests/test_training.py
index f835858..5a4f443 100644
--- a/cellseg_models_pytorch/training/tests/test_training.py
+++ b/cellseg_models_pytorch/training/tests/test_training.py
@@ -9,7 +9,7 @@
 from cellseg_models_pytorch.training.lit import SegmentationExperiment
 
 
-def test_training(img_patch_dir, mask_patch_dir, auto_lr):
+def test_training(img_patch_dir, mask_patch_dir):
     train_ds = SegmentationFolderDataset(
         path=img_patch_dir.as_posix(),
         mask_path=mask_patch_dir.as_posix(),