Repository: huggingface/transformers
Pull request: Added guards against device mismatch errors in some LogitsProcessor classes

There are a few `LogitsProcessor` classes that have a `torch.Tensor` attribute, and when `__call__` is called that tensor may not be on the same device as `scores`. This aims to fix that, not by copying the attribute every time `__call__` happen, but by re-assigning it to `scores` device each time. That way, if the `scores` stay on the same device (which seems likely?) then this becomes a no-op.  

cc. @mht-sharma

Changed files:
filename: src/transformers/generation/logits_process.py
patch: @@ -147,6 +147,7 @@ def __init__(self, min_length: int, eos_token_id: Union[int, List[int], torch.Te
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.eos_token_id = self.eos_token_id.to(scores.device)
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
         eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id)
         scores_processed = scores.clone()
@@ -217,6 +218,7 @@ def __init__(
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.eos_token_id = self.eos_token_id.to(scores.device)
         new_tokens_length = input_ids.shape[-1] - self.prompt_length_to_skip
         scores_processed = scores.clone()
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
@@ -849,6 +851,7 @@ def __init__(
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.epsilon = self.epsilon.to(scores.device)
         probabilities = scores.softmax(dim=-1)
         entropy = torch.distributions.Categorical(logits=scores).entropy()
         eta = torch.min(self.epsilon, torch.sqrt(self.epsilon) * torch.exp(-entropy))[..., None]
@@ -1626,6 +1629,7 @@ def __init__(self, max_length: int, eos_token_id: Union[int, List[int], torch.Te
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.eos_token_id = self.eos_token_id.to(scores.device)
         cur_len = input_ids.shape[-1]
         scores_processed = scores
         if cur_len == self.max_length - 1:
@@ -1837,6 +1841,7 @@ def set_begin_index(self, begin_index):
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.begin_suppress_tokens = self.begin_suppress_tokens.to(scores.device)
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
         suppress_token_mask = isin_mps_friendly(vocab_tensor, self.begin_suppress_tokens)
         scores_processed = scores
@@ -1880,6 +1885,7 @@ def __init__(self, suppress_tokens, device: str = "cpu"):
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.suppress_tokens = self.suppress_tokens.to(scores.device)
         vocab_tensor = torch.arange(scores.shape[-1], device=scores.device)
         suppress_token_mask = isin_mps_friendly(vocab_tensor, self.suppress_tokens)
         scores = torch.where(suppress_token_mask, -float("inf"), scores)
@@ -2343,6 +2349,7 @@ def __init__(self, eos_token_id: Union[int, List[int], torch.Tensor], min_eos_p:
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.eos_token_id = self.eos_token_id.to(scores.device)
         scores_processed = scores
         if self.min_eos_p:
             probs = torch.nn.functional.softmax(scores.float(), dim=-1)