check for None forced tokens (#21793)

src/transformers/generation/flax_logits_process.py

@@ -328,6 +328,7 @@ class FlaxForceTokensLogitsProcessor(FlaxLogitsProcessor):
         # Indexes without forced tokens will have a negative value.
         force_token_array = jnp.ones((max(force_token_map.keys()) + 1), dtype=jnp.int32) * -1
         for index, token in force_token_map.items():
+            if token is not None:
                 force_token_array = force_token_array.at[index].set(token)
         self.force_token_array = jnp.int32(force_token_array)