fix(decoder): only use available cores

huggingface · Jan 31, 2024 · dceed18 · dceed18
1 parent 4531aa2
commit dceed18
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py
@@ -238,7 +238,7 @@ def get_export_config(
             sequence_length = config.max_position_embeddings
         if num_cores is None:
             # Use all available cores
-            num_cores = len(os.listdir("/sys/class/neuron_device/")) * 2
+            num_cores = get_available_cores()
         if auto_cast_type is None:
             auto_cast_type = "fp32"
             if config.torch_dtype == "float16":