bug fixes (#732)

deepjavalibrary · May 19, 2023 · 46ed6c0 · 46ed6c0
1 parent 0007df5
commit 46ed6c0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/deepspeed.py b/engines/python/setup/djl_python/deepspeed.py
@@ -298,7 +298,7 @@ def inference(self, inputs: Input):
                     return_tensors="pt").to(torch.cuda.current_device())
                 with torch.no_grad():
                     output_tokens = self.model.generate(
-                        input_id=tokenized_inputs.input_id,
+                        input_ids=tokenized_inputs.input_ids,
                         attention_mask=tokenized_inputs.attention_mask,
                         **model_kwargs)
                 generated_text = self.tokenizer.batch_decode(

diff --git a/engines/python/setup/djl_python/huggingface.py b/engines/python/setup/djl_python/huggingface.py
@@ -241,7 +241,7 @@ def wrapped_pipeline(inputs, *args, **kwargs):
             with torch.no_grad():
                 output_tokens = model.generate(
                     *args,
-                    input_id=input_tokens.input_id,
+                    input_ids=input_tokens.input_ids,
                     attention_mask=input_tokens.attention_mask,
                     **kwargs)
             generated_text = tokenizer.batch_decode(output_tokens,