enh: run set_num_interop_threads(1) for torchmpo segmenter

DC-analysis · Jun 14, 2024 · e517464 · e517464
1 parent 002003e
commit e517464
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/dcnum/segm/segm_torch/segm_torch_mpo.py b/src/dcnum/segm/segm_torch/segm_torch_mpo.py
@@ -36,8 +36,10 @@ def segment_algorithm(image, *,
         # Set number of pytorch threads to 1, because dcnum is doing
         # all the multiprocessing.
         # https://pytorch.org/docs/stable/generated/torch.set_num_threads.html#torch.set_num_threads
-        torch.set_num_threads(1)
-        torch.set_num_interop_threads(1)
+        if torch.get_num_threads() != 1:
+            torch.set_num_threads(1)
+        if torch.get_num_interop_threads() != 1:
+            torch.set_num_interop_threads(1)
         device = torch.device("cpu")
 
         # Load model and metadata