From 74a094ef8a8c0d350cffac90fcbdbfb6c235efcf Mon Sep 17 00:00:00 2001 From: zhijianma Date: Tue, 14 Nov 2023 20:16:45 +0800 Subject: [PATCH] fixz: set torch num threads to 1 --- data_juicer/ops/filter/clip_similarity_filter.py | 2 +- tests/ops/filter/test_clip_similarity_filter.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/data_juicer/ops/filter/clip_similarity_filter.py b/data_juicer/ops/filter/clip_similarity_filter.py index 8b7d90ee5..81356386e 100644 --- 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 index 12fe318e1..270cb6a4c 100644 --- 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__':