fix: add param to constructor

sileod · Sep 19, 2023 · 48d8e1a · 48d8e1a
1 parent 4077d51
commit 48d8e1a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/tasknet/utils.py b/src/tasknet/utils.py
@@ -214,6 +214,9 @@ def load_pipeline(
             ] += adapter.Z[task_index]
 
     pipe = TextClassificationPipeline(
-        model=model, tokenizer=tokenizer, device=device, return_all_scores=True
+        model=model,
+        tokenizer=tokenizer,
+        device=device,
+        return_all_scores=return_all_scores,
     )
     return pipe