Skip to content

Commit

Permalink
scripts/vsmlrt.py: add param num_threads to the OV_CPU backend
Browse files Browse the repository at this point in the history
  • Loading branch information
WolframRhodium committed Feb 7, 2023
1 parent 191dfda commit 508865d
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion scripts/vsmlrt.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
__version__ = "3.15.8"
__version__ = "3.15.9"

__all__ = [
"Backend", "BackendV2",
Expand Down Expand Up @@ -96,6 +96,7 @@ class OV_CPU:
bind_thread: bool = True
fp16_blacklist_ops: typing.Optional[typing.Sequence[str]] = None
bf16: bool = False
num_threads: int = 0

# internal backend attributes
supports_onnx_serialization: bool = True
Expand Down Expand Up @@ -1323,6 +1324,7 @@ def _inference(
config = lambda: dict(
CPU_THROUGHPUT_STREAMS=backend.num_streams,
CPU_BIND_THREAD="YES" if backend.bind_thread else "NO",
CPU_THREADS_NUM=backend.num_threads,
ENFORCE_BF16="YES" if backend.bf16 else "NO"
)

Expand Down Expand Up @@ -1570,12 +1572,14 @@ def OV_CPU(*,
num_streams: typing.Union[int, str] = 1,
bf16: bool = False,
bind_thread: bool = True,
num_threads: int = 0,
**kwargs
) -> Backend.OV_CPU:
return Backend.OV_CPU(
num_streams=num_streams,
bf16=bf16,
bind_thread=bind_thread,
num_threads=num_threads,
**kwargs
)

Expand Down

0 comments on commit 508865d

Please sign in to comment.