Skip to content

Commit

Permalink
[bugfix] qa tests - KeyError: 'num_cores' (#1152)
Browse files Browse the repository at this point in the history
* [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
  • Loading branch information
bfineran committed Jul 28, 2023
1 parent c988fe7 commit 2c9035e
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/deepsparse/base_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,10 @@ def create(
"Overriding input shapes not supported with Bucketing enabled"
)
if not kwargs.get("context", None):
context = Context(num_cores=kwargs["num_cores"])
context = Context(
num_cores=kwargs.get("num_cores"),
num_streams=kwargs.get("num_streams"),
)
kwargs["context"] = context
buckets = pipeline_constructor.create_pipeline_buckets(
task=task,
Expand Down

0 comments on commit 2c9035e

Please sign in to comment.