fix encoding for inf2 (#534)

* fix encoding * add comments
deepjavalibrary · Mar 8, 2023 · 583dbc0 · 583dbc0
1 parent 9ae1e0c
commit 583dbc0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/transformer-neuronx.py b/engines/python/setup/djl_python/transformer-neuronx.py
@@ -151,8 +151,9 @@ def infer(self, inputs):
                     f"{self.batch_size} batch size not equal to {len(input_text)} prompt size"
                 )
             with torch.inference_mode():
-                input_ids = torch.as_tensor(
-                    [self.tokenizer.encode(text) for text in input_text])
+                # inf 2 needs padding
+                input_ids = self.tokenizer.batch_encode_plus(
+                    input_text, return_tensors="pt", padding=True)['input_ids']
                 generated_sequence = self.model.sample(
                     input_ids, sequence_length=seq_length)
                 result = [