when l use the model generate method within @tf.function, it encounterd a mistake
Author: HelloWorldU
Pull request: no
Labels: name: bug

### System Info

transformers version: 4.43.3
Platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.17
Python version: 3.8.19
Huggingface_hub version: 0.24.5
Safetensors version: 0.4.4
Accelerate version: not installed
Accelerate config: not found
PyTorch version (GPU?): not installed (NA)
Tensorflow version (GPU?): 2.7.0 (True)
Flax version (CPU?/GPU?/TPU?): not installed (NA)
Jax version: not installed
JaxLib version: not installed
Using distributed or parallel set-up in script?: Yes, using TensorFlow MirroredStrategy for distributed training.

### Who can help?

@ArthurZucker  the original issues is #33329, thanks a lot

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

''
def train_generator_step(self, input_ids, attention_mask, labels, styles, max_len, step, accumulation_steps=4,
lambda_rec=1.0, lambda_lm=1.0, lambda_adv=1.0, lambda_kl=1.0, gamma=1.0):

    max_len = tf.constant(max_len, dtype=tf.int32)
    max_len_value = max_len
    seq_len = input_ids.shape[2]
    max_new_tokens = tf.maximum(max_len_value - seq_len - 10, 1)
    max_new_tokens = tf.cast(max_new_tokens, tf.int32)
    max_new_tokens = tf.maximum(max_new_tokens, 1)
    max_new_tokens = tf.constant(max_new_tokens, dtype=tf.int32)

    @tf.function
    def step_fn(input_ids=input_ids, attention_mask=attention_mask, labels=labels, styles=styles, accumulation_steps=accumulation_steps, 
                lambda_rec=lambda_rec, lambda_lm=lambda_lm, lambda_adv=lambda_adv, lambda_kl=lambda_kl, gamma=gamma):
        with tf.GradientTape() as tape:
            tf.debugging.enable_check_numerics()
            
            accumulation_steps, lambda_rec, lambda_lm, lambda_adv, lambda_kl, gamma = pr.conv_tensor_to_float(accumulation_steps, lambda_rec, lambda_lm, lambda_adv, lambda_kl, gamma)

            epsilon = 1e-6 # 快速修复

            """
            we firstly to reshape the input
            """
            actual_shape = tf.shape(input_ids)
            input_ids = tf.reshape(input_ids, (actual_shape[0] * actual_shape[1], actual_shape[2]))
            attention_mask = tf.reshape(attention_mask, (actual_shape[0] * actual_shape[1], actual_shape[2]))

            """
            then, we repeat styles and labels
            """
            styles = tf.repeat(styles, repeats=actual_shape[0])
            labels = tf.repeat(labels, repeats=actual_shape[0], axis=0)

            # 嵌入风格标签
            style_embeddings = self.embedding(styles) # [num_devices * batch_size, n_embd]
            print("Style embeddings shape:", style_embeddings.shape)  # Debug info
            
            # 将输入 ID 嵌入到相同的嵌入空间
            input_embeddings = self.gen.transformer.wte(input_ids) # [num_devices * batch_size, seq_len, n_embd]
            print("Input embeddings shape:", input_embeddings.shape)  # Debug info
            
            extended_input_embeddings = input_embeddings + tf.expand_dims(style_embeddings, axis=1)
            print("Extended embeddings shape:", extended_input_embeddings.shape)  # Debug info

            input_ids, attention_mask, labels, styles = dis.convert_tensor(input_ids, attention_mask, labels, styles)

            outputs = self.gen(input_ids=input_ids, attention_mask=attention_mask, training=True)
            logits = outputs.logits
            print("Logits shape:", logits.shape)  # Debug info
            print("Logits dtype:", logits.dtype)  # Debug info
            print("labels shape:", labels.shape)  # Debug info

            loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
            mask = tf.cast(labels != -100, logits.dtype)
            print("Mask shape:", mask.shape)  # Debug info
            print("Mask dtype:", mask.dtype)  # Debug info
    
            # Check for NaN or Inf in logits
            tf.debugging.check_numerics(logits, "Logits contain NaN or Inf")

            # 各个损失
            rec_loss = loss_fn(tf.where(labels == -100, tf.zeros_like(labels, dtype=logits.dtype), tf.cast(labels, logits.dtype)), logits)
            rec_loss = tf.reduce_sum(rec_loss * mask) / (tf.reduce_sum(mask) + epsilon)
            rec_loss = tf.cast(rec_loss, tf.float32)
            print("Reconstruction loss:", rec_loss)  # Debug info

            for var in self.gen.trainable_variables:
                tf.debugging.check_numerics(var, message="Model weight check")

            print("Input shapes:", 
                    "input_ids:", input_ids.shape, 
                    "input_ids dtype:", input_ids.dtype,
                    "attention_mask:", attention_mask.shape, 
                    "labels:", labels.shape, 
                    "styles:", styles.shape,
                    "max_len_value:", max_len_value)

            new_shape = tf.shape(input_ids)
            print("New shape:", new_shape)  # Debug info
            print("Seq len:", seq_len)  # Debug info
            
            # max_new_tokens = tf.maximum(max_new_tokens, 1)
            print("Max length:", max_len_value)  # Debug info
            print("Max new tokens:", max_new_tokens)  # Debug info
            if isinstance(max_new_tokens, tf.Tensor):
                print("Max new tokens:", tf.get_static_value(max_new_tokens))  # Debug info
            batch_size = new_shape[0]

            # 扩展 input_ids
            """
            at here, we need to padding to -> [batch_size, max_new_tokens]
            """
            padding = tf.zeros((batch_size, max_new_tokens), dtype=input_ids.dtype)
            print("Padding shape:", padding.shape)  # Debug info
            
            extended_input_ids = tf.concat([input_ids, padding], axis=1)
            extended_attention_mask1 = tf.concat([attention_mask, tf.zeros((tf.shape(attention_mask)[0],
                                        max_new_tokens), dtype=attention_mask.dtype)], axis=1)
            
            extended_input_ids = tf.cast(extended_input_ids, tf.int32)
            extended_attention_mask1 = tf.cast(extended_attention_mask1, tf.float32)
            
            print(f"Extended input_ids shape: {extended_input_ids.shape}")
            print(f"Extended attention_mask shape: {extended_attention_mask1.shape}")

            # 确保最大长度大于最小长度
            max_length = max_len_value + max_new_tokens
            min_length = 1
            # tf.print("max_length:", max_length)
            # tf.print("min_length:", min_length)

            tf.debugging.assert_greater(max_length, min_length, message=f"max_length ({max_length}) must be greater than min_length ({min_length})")
            
            pad_token_id = int(self.tokenizer.pad_token_id)
            eos_token_id = int(self.tokenizer.eos_token_id)
            bos_token_id = int(self.tokenizer.bos_token_id)

            try:                    
                generated_ids = self.gen.generate(
                    extended_input_ids, 
                    attention_mask=extended_attention_mask1, 
                    max_new_tokens=max_new_tokens,
                    pad_token_id=pad_token_id,
                    eos_token_id=eos_token_id,
                    bos_token_id=bos_token_id,
                    # use_cache=True,
                    # num_beams=1,  # 使用贪婪搜索
                    do_sample=False,  # 不使用采样
                    # temperature=1.0,  # 降低随机性
                )
                print("Generation successful. Generated IDs shape:", generated_ids.shape)

            except Exception as e:
                print(f"Error during generation: {e}")
                print(f"input_ids shape: {input_ids.shape}")
                print(f"attention_mask shape: {attention_mask.shape}")
                print(f"max_len_value: {max_len_value}")
                raise
'''

"""
Relevant Message
"""

Training gen
Style embeddings shape: (4, 768)
Input embeddings shape: (4, 72, 768)
Extended embeddings shape: (4, 72, 768)
Logits shape: (4, 72, 21128)
Logits dtype: <dtype: 'float16'>
labels shape: (4, 72)
Mask shape: (4, 72)
Mask dtype: <dtype: 'float16'>
Reconstruction loss: Tensor("Cast_3:0", shape=(), dtype=float32)
Input shapes: input_ids: (4, 72) input_ids dtype: <dtype: 'int32'> attention_mask: (4, 72) labels: (4, 72) styles: (4,) max_len_value: tf.Tensor(125, shape=(), dtype=int32)
New shape: Tensor("Shape_1:0", shape=(2,), dtype=int32)
Seq len: 72
Max length: tf.Tensor(125, shape=(), dtype=int32)
Max new tokens: tf.Tensor(43, shape=(), dtype=int32)
Max new tokens: 43
Padding shape: (4, 43)
Extended input_ids shape: (4, 115)
Extended attention_mask shape: (4, 115)
/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py:377: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )
return py_builtins.overload_of(f)(*args)
Error during generation: max_new_tokens must be greater than 0, but is 43.
input_ids shape: (4, 72)
attention_mask shape: (4, 72)
max_len_value: 125
Traceback (most recent call last):
File "train.py", line 530, in
train_model.train(train_tf_dataset_X, train_tf_dataset_Y, valid_tf_dataset_X, valid_tf_dataset_Y, trainconfig.epochs)
File "train.py", line 306, in train
rec_loss, lm_loss, adv_loss, kl_loss, current_lr, accuracy, total_gen_loss = self.distributed_train_generator_step(
File "train.py", line 138, in distributed_train_generator_step
loss, rec_loss, lm_loss, adv_loss, kl_loss, current_lr, accuracy = self.strategy.run(
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py", line 1316, in run
return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py", line 2892, in call_for_each_replica
return self._call_for_each_replica(fn, args, kwargs)
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/distribute/mirrored_strategy.py", line 677, in _call_for_each_replica
return mirrored_run.call_for_each_replica(
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/distribute/mirrored_run.py", line 104, in call_for_each_replica
return _call_for_each_replica(strategy, fn, args, kwargs)
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/distribute/mirrored_run.py", line 246, in _call_for_each_replica
coord.join(threads)
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/training/coordinator.py", line 389, in join
six.reraise(*self._exc_info_to_raise)
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/six.py", line 719, in reraise
raise value
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/training/coordinator.py", line 297, in stop_on_exception
yield
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/distribute/mirrored_run.py", line 346, in run
self.main_result = self.main_fn(*self.main_args, **self.main_kwargs)
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py", line 601, in wrapper
return func(*args, **kwargs)
File "train.py", line 133, in generator_step
loss, rec_loss, lm_loss, adv_loss, kl_loss, current_lr, accuracy, gradients = self.model.train_generator_step(*args, **kwargs)
File "/root/autodl-tmp/model/model.py", line 302, in train_generator_step
step_total_loss, step_rec_loss, step_lm_loss, step_adv_loss, step_kl_loss, step_gradients, step_accuracy = step_fn(
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1129, in autograph_handler
raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

File "/root/autodl-tmp/model/model.py", line 220, in step_fn  *
    generated_ids = self.gen.generate(
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/transformers/generation/tf_utils.py", line 738, in generate  *
    model_kwargs = generation_config.update(**kwargs)  # All unused kwargs must be model kwargs
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/transformers/generation/configuration_utils.py", line 1207, in update  *
    self.validate()
File "/root/miniconda3/envs/gpt2-env/lib/python3.8/site-packages/transformers/generation/configuration_utils.py", line 544, in validate  *
    raise ValueError(f"`max_new_tokens` must be greater than 0, but is {self.max_new_tokens}.")

ValueError: `max_new_tokens` must be greater than 0, but is 43.
"""
Definetely, l occured this mistake within @tf.function, and there is no logical mistake when l dubug my code under eager-excution model, similarly, when i use max_length and min_length paramters, it would be occured to "ValueError: max_length must be greater than min_length, 1 is larger than 128.", like this. But, when l set the paramter"max_new_tokens" as a constant value like 50, it would be fine, l donno what leads this, and debug this for at least 20 times.
"""

### Expected behavior

Of course, the value of my variable is dynamic, but I have already defined it outside the graph and used it as a parameter. My expected behavior should be 43 as max_new_token, but it reported an error.

Comments:
sry, the original issus is #33229 

Hey! Thanks for posting! Which model are you using? Also one weird thing is: `train_generator_step` should not use generate, in general using generate is for inference! 

yeah, l know.l use uer/gpt2-chinese-cluecorpussmall model as generator. And l need to use generate to train my GAN, you could find complete code section in my repository, basically, l need the generate id to train my model.

> Hey! Thanks for posting! Which model are you using? Also one weird thing is: `train_generator_step` should not use generate, in general using generate is for inference!



I would recommend not using TF in the mean time 😅 

> I would recommend not using TF in the mean time 😅

anyway, thanks for looking, l aint mean to provoke you

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.