Also set the number of threads

microsoft · Aug 6, 2024 · 2f7b8e3 · 2f7b8e3
1 parent 091dec6
commit 2f7b8e3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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)