enh: run set_num_interop_threads(1) for torchmpo segmenter

DC-analysis · Jun 14, 2024 · 002003e · 002003e
1 parent d7c7d09
commit 002003e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,5 @@
+0.23.4
+ - enh: run set_num_interop_threads(1) for torchmpo segmenter
 0.23.3
  - fix: ignore non-file-type-like basins
  - fix: workaround for slow reading from HDF5 (don't use index arrays)

diff --git a/src/dcnum/segm/segm_torch/segm_torch_mpo.py b/src/dcnum/segm/segm_torch/segm_torch_mpo.py
@@ -37,6 +37,7 @@ def segment_algorithm(image, *,
         # 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)
         device = torch.device("cpu")
 
         # Load model and metadata