Skip to content

Commit 2c9035e

Browse files
authored
[bugfix] qa tests - KeyError: 'num_cores' (#1152)
* [bugfix] qa tests - `KeyError: 'num_cores'` if num_cores is not an explicit kwarg of a bucketing pipeline it will trigger a key error on the changed line. this PR uses `.get` to default the value to `None` (this is the expected default of `Context`) * add num_streams as well
1 parent c988fe7 commit 2c9035e

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

Diff for: src/deepsparse/base_pipeline.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,10 @@ def create(
196196
"Overriding input shapes not supported with Bucketing enabled"
197197
)
198198
if not kwargs.get("context", None):
199-
context = Context(num_cores=kwargs["num_cores"])
199+
context = Context(
200+
num_cores=kwargs.get("num_cores"),
201+
num_streams=kwargs.get("num_streams"),
202+
)
200203
kwargs["context"] = context
201204
buckets = pipeline_constructor.create_pipeline_buckets(
202205
task=task,

0 commit comments

Comments
 (0)