GPT-J finetuning for  Sequential Classification fails
Author: jojivk73
Pull request: no

### System Info

Framework : HF+ TF
Device : CPU Sapphire Rapids
Issue : GPt-J crashes in HF transformer when run in graph mode.Eager mode runs fine.
Any fix/resolution/workaround appreciated.



### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction
```python
from datasets import load_dataset
from transformers import  (
    AutoTokenizer,
    DataCollatorWithPadding,
    TFAutoModelForSequenceClassification,
    TFAutoModelForCausalLM,
    AutoConfig,
)
from tensorflow.keras.losses import SparseCategoricalCrossentropy
import numpy as np
import tensorflow as tf


IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "
"
DEFAULT_BOS_TOKEN = "
"
DEFAULT_UNK_TOKEN = "<unk>"

tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')
tf.config.threading.set_inter_op_parallelism_threads(2)
tf.config.threading.set_intra_op_parallelism_threads(112)

raw_datasets = load_dataset("glue", "cola")
checkpoint = "EleutherAI/gpt-j-6b"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
#print(raw_datasets["train"][0])

label_list = raw_datasets["train"].features["label"].names
num_labels = len(label_list)
config = AutoConfig.from_pretrained(checkpoint, num_labels=num_labels)


special_tokens_dict = dict()
if tokenizer.pad_token is None:
   special_tokens_dict["pad_token"] = DEFAULT_PAD_TOKEN
if tokenizer.eos_token is None:
   special_tokens_dict["eos_token"] = DEFAULT_EOS_TOKEN
if tokenizer.bos_token is None:
   special_tokens_dict["bos_token"] = DEFAULT_BOS_TOKEN
if tokenizer.unk_token is None:
   special_tokens_dict["unk_token"] = DEFAULT_UNK_TOKEN

tokenizer.add_special_tokens(special_tokens_dict)
config.pad_token_id=0
print(config)
def tokenize_function(example):
    return tokenizer(example["sentence"], padding="longest", truncation=True)


tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
print(tokenized_datasets)

#data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False, return_tensors="tf")

tf_train_dataset = tokenized_datasets["train"].to_tf_dataset(
    columns=["input_ids"], #, "token_type_ids"],
    label_cols=["label"],
    shuffle=True,
    collate_fn=data_collator,
    batch_size=64,
)

tf_validation_dataset = tokenized_datasets["validation"].to_tf_dataset(
    #columns=["attention_mask", "input_ids"],
    columns=["input_ids"],
    label_cols=["label"],
    shuffle=False,
    collate_fn=data_collator,
    batch_size=64,
)

model = TFAutoModelForSequenceClassification.from_pretrained(checkpoint, config=config)
model.resize_token_embeddings(len(tokenizer))

from tensorflow.keras.optimizers.schedules import PolynomialDecay

num_epochs=3
# The number of training steps is the number of samples in the dataset, divided by the batch size then multiplied
# by the total number of epochs. Note that the tf_train_dataset here is a batched tf.data.Dataset,
# not the original Hugging Face Dataset, so its len() is already num_samples // batch_size.
num_train_steps = len(tf_train_dataset) * num_epochs
lr_scheduler = PolynomialDecay(
    initial_learning_rate=2e-5, end_learning_rate=0.0, decay_steps=num_train_steps
)
from tensorflow.keras.optimizers import Adam

opt = Adam(learning_rate=lr_scheduler)

model.compile(
    optimizer=opt,
    loss=SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
    #run_eagerly=True,
)
#import pdb
#pdb.set_trace()
model.fit(
    tf_train_dataset,
    validation_data=tf_validation_dataset,
    epochs=3,
)


### Expected behavior

The model fails when run_eagerly is False(default). If run_Eagerly is True, it does not fail.

Failure log
================================================================
    File "/localdisk/jojimonv/tfv2_p39/lib64/python3.9/site-packages/keras/engine/training.py", line 1268, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/localdisk/jojimonv/tfv2_p39/lib64/python3.9/site-packages/keras/engine/training.py", line 1249, in run_step  **
        outputs = model.train_step(data)
    File "/localdisk/jojimonv/tfv2_p39/lib64/python3.9/site-packages/transformers/modeling_tf_utils.py", line 1637, in train_step
        y_pred = self(x, training=True)
    File "/localdisk/jojimonv/tfv2_p39/lib64/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_file05uidiwg.py", line 37, in tf__run_call_with_unpacked_inputs
        retval_ = ag__.converted_call(ag__.ld(func), (ag__.ld(self),), dict(**ag__.ld(unpacked_inputs)), fscope)
    File "/tmp/__autograph_generated_file3qh7zjf_.py", line 57, in tf__call
        ag__.if_stmt(ag__.ld(self).config.pad_token_id is None, if_body_1, else_body_1, get_state_1, set_state_1, ('in_logits', 'sequence_lengths'), 2)
    File "/tmp/__autograph_generated_file3qh7zjf_.py", line 55, in else_body_1
        ag__.if_stmt(ag__.ld(input_ids) is not None, if_body, else_body, get_state, set_state, ('in_logits', 'sequence_lengths'), 2)
    File "/tmp/__autograph_generated_file3qh7zjf_.py", line 47, in if_body
        sequence_lengths = ag__.converted_call(ag__.ld(tf).where, (ag__.ld(sequence_lengths) >= 0, ag__.ld(sequence_lengths), ag__.ld(input_ids).shape[-1] - 1), None, fscope)

    TypeError: Exception encountered when calling layer 'tfgptj_for_sequence_classification' (type TFGPTJForSequenceClassification).
    
    in user code:
    
        File "/localdisk/jojimonv/tfv2_p39/lib64/python3.9/site-packages/transformers/modeling_tf_utils.py", line 834, in run_call_with_unpacked_inputs  *
            return func(self, **unpacked_inputs)
        File "/localdisk/jojimonv/tfv2_p39/lib64/python3.9/site-packages/transformers/models/gptj/modeling_tf_gptj.py", line 875, in call  *
            sequence_lengths = tf.where(sequence_lengths >= 0, sequence_lengths, input_ids.shape[-1] - 1)
    
        TypeError: unsupported operand type(s) for -: 'NoneType' and 'int'
    
    
    Call arguments received by layer 'tfgptj_for_sequence_classification' (type TFGPTJForSequenceClassification):
      • input_ids=tf.Tensor(shape=(None, None), dtype=int64)
      • past_key_values=None
      • attention_mask=None
      • token_type_ids=None
      • position_ids=None
      • head_mask=None
      • inputs_embeds=None
      • labels=None
      • use_cache=None
      • output_attentions=None
      • output_hidden_states=None
      • return_dict=None
      • training=True
```

Comments:
AH yes, this is known, one of our CI caught this! 
I'll open a PR !

Sorry got caught up in a lot of things! TF Is a pain to fix 😅 

cc @Rocketknight1 will probably take you a lot less time than me if you have time to tkae #25751 over! (Have not really seen any complaints so very low priority) 

Added the fix in #25751!

@jojivk73 This should now be fixed! Please install from `main` with `pip install git+https://github.com/huggingface/transformers.git` to try it out, and let me know if you have any other issues!