Make EosTokenCriteria compatible with mps (#30376)

huggingface · Apr 23, 2024 · f8fec6b · f8fec6b
1 parent 745bbfe
commit f8fec6b
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/src/transformers/generation/stopping_criteria.py b/src/transformers/generation/stopping_criteria.py
@@ -146,7 +146,18 @@ def __init__(self, eos_token_id: Union[int, List[int]]):
 
     @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
-        is_done = torch.isin(input_ids[:, -1], 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 = (
+                input_ids[:, -1]
+                .tile(self.eos_token_id.shape[0], 1)
+                .eq(self.eos_token_id.unsqueeze(1).to(input_ids.device))
+                .sum(dim=0)
+                .bool()
+                .squeeze()
+            )
+        else:
+            is_done = torch.isin(input_ids[:, -1], self.eos_token_id.to(input_ids.device))
         return is_done