src/transformers/generation/flax_logits_process.py/FlaxSuppressTokensLogitsProcessor/__call__
class FlaxSuppressTokensLogitsProcessor:    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        scores = scores.at[..., self.suppress_tokens].set(-float("inf"))

        return scores