From c62932f29ec1d3a94abbe9dd63852bdde89aa4ab Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Fri, 26 Jan 2024 19:40:14 -0500
Subject: [PATCH] set threads

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/tests/pt/__init__.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/source/tests/pt/__init__.py b/source/tests/pt/__init__.py
index 6ceb116d85..fdbdd73f79 100644
--- a/source/tests/pt/__init__.py
+++ b/source/tests/pt/__init__.py
@@ -1 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import torch
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)