From 2f7b8e309ed2a238fa58540a1064b72e10e6fa6d Mon Sep 17 00:00:00 2001 From: Wessel Bruinsma Date: Tue, 6 Aug 2024 16:46:50 +0200 Subject: [PATCH] Also set the number of threads --- tests/test_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_model.py b/tests/test_model.py index 61dc78f..850af9b 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -80,6 +80,7 @@ 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)