diff --git a/tests/algorithms/test_algorithm_resumption.py b/tests/algorithms/test_algorithm_resumption.py index 6d6145cf62..ccf3e7b2b6 100644 --- a/tests/algorithms/test_algorithm_resumption.py +++ b/tests/algorithms/test_algorithm_resumption.py @@ -53,7 +53,7 @@ def test_algorithm_resumption( copied_model = copy.deepcopy(model) # copy the model so the params will start from the same point - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1) # Reduce training duration and data @@ -82,7 +82,7 @@ def test_algorithm_resumption( # create second trainer, load an intermediate checkpoint # and continue training - optimizer = torch.optim.SGD(copied_model.parameters(), lr=0.1) + optimizer = torch.optim.Adam(copied_model.parameters(), lr=0.1) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1) alg = alg_cls(**alg_kwargs) diff --git a/tests/test_full_nlp.py b/tests/test_full_nlp.py index a3f342500b..7afb8cc08c 100644 --- a/tests/test_full_nlp.py +++ b/tests/test_full_nlp.py @@ -60,7 +60,7 @@ def pretraining_test_helper(tokenizer, model, algorithms, tmp_path, device): model=pretraining_model_copy, train_dataloader=pretraining_train_dataloader, save_folder=str(tmp_path / 'pretraining_checkpoints'), - max_duration='1ba', + max_duration='2ba', seed=17, algorithms=algorithms, device=device, @@ -138,7 +138,7 @@ def finetuning_test_helper(tokenizer, model, algorithms, checkpoint_path, pretra load_weights_only=True, load_strict_model_weights=False, loggers=[rud], - max_duration='1ba', + max_duration='2ba', seed=17, algorithms=algorithms, device=device,