Skip to content

Commit

Permalink
FIX | Update serve_trt.py
Browse files Browse the repository at this point in the history
Remove lines that handle out-of-range IDs which are not needed now and may cause issues.
  • Loading branch information
yinglin-nv authored Jun 27, 2024
1 parent b23cd76 commit 2e50a9b
Showing 1 changed file with 0 additions and 3 deletions.
3 changes: 0 additions & 3 deletions nemo_skills/inference/server/serve_trt.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,9 +160,6 @@ def get_output_single(output_ids, input_length, max_output_len, tokenizer, eos_t
if len(eos_ids) > 0:
outputs = outputs[: eos_ids[0]]
outputs = outputs.tolist()
# somehow sometimes it produces tokens out of range..
outputs = [elem if elem < tokenizer.vocab_size else tokenizer.vocab_size - 1 for elem in outputs]
outputs = [elem if 0 <= elem else 0 for elem in outputs]
return tokenizer.decode(outputs)


Expand Down

0 comments on commit 2e50a9b

Please sign in to comment.