Fix input_data and device order for streaming (deepjavalibrary#809)

KexinFeng · Aug 16, 2023 · 59541c3 · 59541c3
1 parent 2328b6c
commit 59541c3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/deepspeed.py b/engines/python/setup/djl_python/deepspeed.py
@@ -324,8 +324,8 @@ def inference(self, inputs: Input):
                         "DeepSpeed")
                     device = torch.cuda.current_device()
                     outputs.add_stream_content(
-                        stream_generator(self.model, self.tokenizer, device,
-                                         input_data, **model_kwargs))
+                        stream_generator(self.model, self.tokenizer, input_data,
+                                         device, **model_kwargs))
                 return outputs
             if self.task == "text-generation":
                 tokenized_inputs = self.tokenizer(