Revoke cpu_threads in test_transformers_wav2vec2 similar to test_transformers_wav2vec2bert

3manifold · 3manifold · commit 783d03b46fd7 · 2026-02-04T14:10:07.000+01:00
diff --git a/python/tests/test_transformers.py b/python/tests/test_transformers.py
@@ -1011,13 +1011,13 @@ def test_transformers_wav2vec2(
         )
 
         device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
-        # cpu_threads = int(os.environ.get("OMP_NUM_THREADS", 0))
+        cpu_threads = int(os.environ.get("OMP_NUM_THREADS", 0))
         model = ctranslate2.models.Wav2Vec2(
             output_dir,
             device=device,
             device_index=[0],
             compute_type="int8",
-            intra_threads=1,
+            intra_threads=cpu_threads,
             inter_threads=1,
         )
 

Original file line number	Diff line number	Diff line change
`@@ -1011,13 +1011,13 @@ def test_transformers_wav2vec2(`
`1011`	`1011`	`)`
`1012`	`1012`
`1013`	`1013`	`device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"`
`1014`		`- # cpu_threads = int(os.environ.get("OMP_NUM_THREADS", 0))`
	`1014`	`+ cpu_threads = int(os.environ.get("OMP_NUM_THREADS", 0))`
`1015`	`1015`	`model = ctranslate2.models.Wav2Vec2(`
`1016`	`1016`	`output_dir,`
`1017`	`1017`	`device=device,`
`1018`	`1018`	`device_index=[0],`
`1019`	`1019`	`compute_type="int8",`
`1020`		`- intra_threads=1,`
	`1020`	`+ intra_threads=cpu_threads,`
`1021`	`1021`	`inter_threads=1,`
`1022`	`1022`	`)`
`1023`	`1023`