Make EosTokenCriteria compatible with mps (#30376)

src/transformers/generation/stopping_criteria.py

@@ -481,7 +481,18 @@ class EosTokenCriteria(StoppingCriteria):
 
     @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