awslabs / sockeye

Sequence-to-sequence framework with a focus on Neural Machine Translation based on PyTorch
https://awslabs.github.io/sockeye/
Apache License 2.0
1.21k stars 324 forks source link

Two fixes to SampleK #1086

Closed tdomhan closed 1 year ago

tdomhan commented 1 year ago

Two small fixes to SampleK. Before the device was not set correctly leading to issues when running sampling on GPUs. Furthermore, SampleK did not return the top-k values correctly.

Pull Request Checklist

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.