add num_streams as well

neuralmagic · Jul 28, 2023 · 535069b · 535069b
1 parent 93bae7b
commit 535069b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/deepsparse/base_pipeline.py b/src/deepsparse/base_pipeline.py
@@ -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.get("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,