fixz: set torch num threads to 1

modelscope · Nov 14, 2023 · 74a094e · 74a094e
1 parent 8288a4a
commit 74a094e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/data_juicer/ops/filter/clip_similarity_filter.py b/data_juicer/ops/filter/clip_similarity_filter.py
@@ -10,7 +10,7 @@
 from ..op_fusion import LOADED_IMAGES
 
 # avoid hanging when calling clip in multiprocessing
-torch.get_num_threads()
+torch.set_num_threads(1)
 
 
 @OPERATORS.register_module('clip_similarity_filter')

diff --git a/tests/ops/filter/test_clip_similarity_filter.py b/tests/ops/filter/test_clip_similarity_filter.py
@@ -202,7 +202,7 @@ def test_multi_process(self):
                                   any_or_all='any',
                                   min_ratio=0.2,
                                   max_ratio=0.9)
-        self._run_filter(dataset, tgt_list, op, num_proc=1)
+        self._run_filter(dataset, tgt_list, op, num_proc=4)
 
 
 if __name__ == '__main__':