From 0103d00c906b11440603639049cfa82e3902f098 Mon Sep 17 00:00:00 2001 From: Wessel Bruinsma Date: Tue, 6 Aug 2024 16:59:05 +0200 Subject: [PATCH] Add seed at the beginning --- tests/test_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_model.py b/tests/test_model.py index 850af9b..e1388c7 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -11,6 +11,7 @@ from aurora import AuroraSmall, Batch, Metadata +torch.manual_seed(0) torch.use_deterministic_algorithms(True) @@ -80,7 +81,6 @@ def test_aurora_small() -> None: # Load the checkpoint and run the model. model.load_checkpoint(os.environ["HUGGINGFACE_REPO"], "aurora-0.25-small-pretrained.ckpt") torch.manual_seed(0) # Very important to seed! The test data was generated using this. - torch.set_num_threads(1) with torch.inference_mode(): pred = model.forward(batch)