Generate: get the correct beam index on eos token (#18851)

2022-09-05 19:35:47 +01:00
parent c6d3daba54
commit d4dbd7ca59
2 changed files with 2 additions and 2 deletions
--- a/tests/generation/test_generation_beam_search.py
+++ b/tests/generation/test_generation_beam_search.py
@@ -172,7 +172,7 @@ class BeamSearchTester:
                input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][1].tolist()
            )
            self.parent.assertListEqual(
-                expected_beam_indices + [next_indices[batch_idx, 1].item()],
+                expected_beam_indices + [correct_idx],
                torch.tensor(beam_scorer._beam_hyps[batch_idx].beams[0][2]).tolist(),
            )