You may meet "RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)" in beam_search.k_best_ouputs Function

[JerryZhouBUPT]

You can add outputs = outputs.to(device) at this place to fix this bug.

[JerryZhouBUPT]

Additionally, you need to pass the device argument into this function.