fix EosTokenCriteria

2025-10-20 17:13:56 +08:00 · 2024-05-20 10:43:14 +02:00
1 changed files with 2 additions and 1 deletions
--- a/src/transformers/generation/stopping_criteria.py
+++ b/src/transformers/generation/stopping_criteria.py
@ -481,6 +481,7 @@ class EosTokenCriteria(StoppingCriteria):

    @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+        self.eos_token_id = self.eos_token_id.to(input_ids.device)
        if input_ids.device.type == "mps":
            # https://github.com/pytorch/pytorch/issues/77764#issuecomment-2067838075
            is_done = (
@ -492,7 +493,7 @@ class EosTokenCriteria(StoppingCriteria):
                .squeeze()
            )
        else:
-            is_done = torch.isin(input_ids[:, -1], self.eos_token_id.to(input_ids.device))
+            is_done = torch.isin(input_ids[:, -1], self.eos_token_id)
        return is_done