Fix bug in llm_finetuning_test.py

Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Nov 28, 2024 · 0c40eb1 · 0c40eb1
1 parent 284011c
commit 0c40eb1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/project/algorithms/llm_finetuning_test.py b/project/algorithms/llm_finetuning_test.py
@@ -122,13 +122,15 @@ def test_initialization_is_reproducible(
         seed: int,
         tensor_regression: TensorRegressionFixture,
         trainer: lightning.Trainer,
+        device: torch.device,
     ):
         super().test_initialization_is_reproducible(
             experiment_config=experiment_config,
             datamodule=datamodule,
             seed=seed,
             tensor_regression=tensor_regression,
             trainer=trainer,
+            device=device,
         )
 
     @pytest.mark.xfail(