T5Tokenizer, TFT5ForConditionalGeneration Graph execution error using tfa.metrics.CohenKappa
### System Info

- `transformers` version: 4.24.0
- Platform: macOS-12.5-arm64-arm-64bit
- Python version: 3.10.9
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.12.1 (False)
- Tensorflow version (GPU?): 2.12.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

Hello Huggingface Team,

I am currently using the T5 modules: T5Tokenizer, TFT5ForConditionalGeneration with the metric tfa.metrics.CohenKappa from Tensorflow Add-ons library. I have fine-tuned a model to achieve multi-class classification which works when I use a different metric, for example, 'accuracy'. The issue is when I exchange the metric to use CohenKappa I get the error found below. If more information is required please let me know. Thank you in advance!

Error:
> InvalidArgumentError                      Traceback (most recent call last)
> Cell In[80], line 15
>       6 #model.fit(tokenized_train_data, validation_data=val_dataset, epochs=num_epochs)
>       8 model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
>       9     filepath='/',
>      10     save_weights_only=True,
>      11     monitor='val_accuracy',
>      12     mode='max',
>      13     save_best_only=True)
> ---> 15 history = t5_model.fit(tokenized_train_data,
>      16              validation_data=tokenized_test_data,
>      17              callbacks=None,#[model_checkpoint_callback],
>      18              batch_size=batch_size,
>      19              epochs=num_epochs)
> 
> File ~/anaconda3/lib/python3.10/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
>      67     filtered_tb = _process_traceback_frames(e.__traceback__)
>      68     # To get the full stack trace, call:
>      69     # `tf.debugging.disable_traceback_filtering()`
> ---> 70     raise e.with_traceback(filtered_tb) from None
>      71 finally:
>      72     del filtered_tb
> 
> File ~/anaconda3/lib/python3.10/site-packages/tensorflow/python/eager/execute.py:52, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
>      50 try:
>      51   ctx.ensure_initialized()
> ---> 52   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
>      53                                       inputs, attrs, num_outputs)
>      54 except core._NotOkStatusException as e:
>      55   if name is not None:
> 
> InvalidArgumentError: Graph execution error:
> 
> Detected at node 'confusion_matrix/assert_less/Assert/AssertGuard/Assert' defined at (most recent call last):
>     File "/Users/pj/anaconda3/lib/python3.10/runpy.py", line 196, in _run_module_as_main
>       return _run_code(code, main_globals, None,
>     File "/Users/pj/anaconda3/lib/python3.10/runpy.py", line 86, in _run_code
>       exec(code, run_globals)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel_launcher.py", line 17, in <module>
>       app.launch_new_instance()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/traitlets/config/application.py", line 992, in launch_instance
>       app.start()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelapp.py", line 711, in start
>       self.io_loop.start()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tornado/platform/asyncio.py", line 199, in start
>       self.asyncio_loop.run_forever()
>     File "/Users/pj/anaconda3/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
>       self._run_once()
>     File "/Users/pj/anaconda3/lib/python3.10/asyncio/base_events.py", line 1906, in _run_once
>       handle._run()
>     File "/Users/pj/anaconda3/lib/python3.10/asyncio/events.py", line 80, in _run
>       self._context.run(self._callback, *self._args)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 510, in dispatch_queue
>       await self.process_one()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 499, in process_one
>       await dispatch(*args)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 406, in dispatch_shell
>       await result
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 729, in execute_request
>       reply_content = await reply_content
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 411, in do_execute
>       res = shell.run_cell(
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/zmqshell.py", line 531, in run_cell
>       return super().run_cell(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 2961, in run_cell
>       result = self._run_cell(
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3016, in _run_cell
>       result = runner(coro)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
>       coro.send(None)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3221, in run_cell_async
>       has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3400, in run_ast_nodes
>       if await self.run_code(code, result, async_=asy):
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3460, in run_code
>       exec(code_obj, self.user_global_ns, self.user_ns)
>     File "/var/folders/wv/10kjqk217c5039dg4pbqggh00000gn/T/ipykernel_39051/3646968355.py", line 15, in <module>
>       history = t5_model.fit(tokenized_train_data,
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
>       return fn(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1685, in fit
>       tmp_logs = self.train_function(iterator)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1284, in train_function
>       return step_function(self, iterator)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1268, in step_function
>       outputs = model.distribute_strategy.run(run_step, args=(data,))
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in run_step
>       outputs = model.train_step(data)
>     File "/var/folders/wv/10kjqk217c5039dg4pbqggh00000gn/T/ipykernel_39051/2458810848.py", line 36, in train_step
>       self.compiled_metrics.update_state(labels, tf.argmax(outputs.logits, -1))
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/compile_utils.py", line 605, in update_state
>       metric_obj.update_state(y_t, y_p, sample_weight=mask)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/utils/metrics_utils.py", line 77, in decorated
>       update_op = update_state_fn(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/metrics/base_metric.py", line 140, in update_state_fn
>       return ag_update_state(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tensorflow_addons/metrics/cohens_kappa.py", line 150, in update_state
>       return self._update(y_true, y_pred, sample_weight)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tensorflow_addons/metrics/cohens_kappa.py", line 165, in _update_multi_class_model
>       return self._update_confusion_matrix(y_true, y_pred, sample_weight)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tensorflow_addons/metrics/cohens_kappa.py", line 193, in _update_confusion_matrix
>       new_conf_mtx = tf.math.confusion_matrix(
> Node: 'confusion_matrix/assert_less/Assert/AssertGuard/Assert'
> Detected at node 'confusion_matrix/assert_less/Assert/AssertGuard/Assert' defined at (most recent call last):
>     File "/Users/pj/anaconda3/lib/python3.10/runpy.py", line 196, in _run_module_as_main
>       return _run_code(code, main_globals, None,
>     File "/Users/pj/anaconda3/lib/python3.10/runpy.py", line 86, in _run_code
>       exec(code, run_globals)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel_launcher.py", line 17, in <module>
>       app.launch_new_instance()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/traitlets/config/application.py", line 992, in launch_instance
>       app.start()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelapp.py", line 711, in start
>       self.io_loop.start()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tornado/platform/asyncio.py", line 199, in start
>       self.asyncio_loop.run_forever()
>     File "/Users/pj/anaconda3/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
>       self._run_once()
>     File "/Users/pj/anaconda3/lib/python3.10/asyncio/base_events.py", line 1906, in _run_once
>       handle._run()
>     File "/Users/pj/anaconda3/lib/python3.10/asyncio/events.py", line 80, in _run
>       self._context.run(self._callback, *self._args)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 510, in dispatch_queue
>       await self.process_one()
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 499, in process_one
>       await dispatch(*args)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 406, in dispatch_shell
>       await result
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 729, in execute_request
>       reply_content = await reply_content
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 411, in do_execute
>       res = shell.run_cell(
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/ipykernel/zmqshell.py", line 531, in run_cell
>       return super().run_cell(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 2961, in run_cell
>       result = self._run_cell(
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3016, in _run_cell
>       result = runner(coro)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
>       coro.send(None)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3221, in run_cell_async
>       has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3400, in run_ast_nodes
>       if await self.run_code(code, result, async_=asy):
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3460, in run_code
>       exec(code_obj, self.user_global_ns, self.user_ns)
>     File "/var/folders/wv/10kjqk217c5039dg4pbqggh00000gn/T/ipykernel_39051/3646968355.py", line 15, in <module>
>       history = t5_model.fit(tokenized_train_data,
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
>       return fn(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1685, in fit
>       tmp_logs = self.train_function(iterator)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1284, in train_function
>       return step_function(self, iterator)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1268, in step_function
>       outputs = model.distribute_strategy.run(run_step, args=(data,))
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in run_step
>       outputs = model.train_step(data)
>     File "/var/folders/wv/10kjqk217c5039dg4pbqggh00000gn/T/ipykernel_39051/2458810848.py", line 36, in train_step
>       self.compiled_metrics.update_state(labels, tf.argmax(outputs.logits, -1))
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/engine/compile_utils.py", line 605, in update_state
>       metric_obj.update_state(y_t, y_p, sample_weight=mask)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/utils/metrics_utils.py", line 77, in decorated
>       update_op = update_state_fn(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/keras/metrics/base_metric.py", line 140, in update_state_fn
>       return ag_update_state(*args, **kwargs)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tensorflow_addons/metrics/cohens_kappa.py", line 150, in update_state
>       return self._update(y_true, y_pred, sample_weight)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tensorflow_addons/metrics/cohens_kappa.py", line 165, in _update_multi_class_model
>       return self._update_confusion_matrix(y_true, y_pred, sample_weight)
>     File "/Users/pj/anaconda3/lib/python3.10/site-packages/tensorflow_addons/metrics/cohens_kappa.py", line 193, in _update_confusion_matrix
>       new_conf_mtx = tf.math.confusion_matrix(
> Node: 'confusion_matrix/assert_less/Assert/AssertGuard/Assert'
> 2 root error(s) found.
>   (0) INVALID_ARGUMENT:  assertion failed: [`labels` out of bound] [Condition x < y did not hold element-wise:] [x (confusion_matrix/control_dependency:0) = ] [[220 1][209...]...] [y (confusion_matrix/Cast:0) = ] [4]
> 	 [[{{node confusion_matrix/assert_less/Assert/AssertGuard/Assert}}]]
> 	 [[gradient_tape/tft5_for_conditional_generation/decoder/block_._5/layer_._0/SelfAttention/mul_1/_596]]
>   (1) INVALID_ARGUMENT:  assertion failed: [`labels` out of bound] [Condition x < y did not hold element-wise:] [x (confusion_matrix/control_dependency:0) = ] [[220 1][209...]...] [y (confusion_matrix/Cast:0) = ] [4]
> 	 [[{{node confusion_matrix/assert_less/Assert/AssertGuard/Assert}}]]
> 0 successful operations.
> 0 derived errors ignored. [Op:__inference_train_function_393119]

### Who can help?

@Rocketknight1 @gante  @ArthurZucker

### 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_step(self, inputs):
  input_ids = inputs['input_ids']
  attention_mask = inputs['attention_mask']
  labels = inputs['labels']
  labels_mask = inputs['labels_mask']
  with tf.GradientTape() as tape:
    outputs = self(input_ids=input_ids,
                       attention_mask=attention_mask,
                       labels=labels,
                       decoder_attention_mask=labels_mask,
                       training=True
                  )
    loss = self.compiled_loss(labels, outputs.logits, regularization_losses=self.losses)
  self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
  self.compiled_metrics.update_state(labels, outputs.logits) ## error happens here
  return_metrics = {}
  for metric in self.metrics:
    result = metric.result()
    if isinstance(result, dict):
        return_metrics.update(result)
    else:
        return_metrics[metric.name] = result
  if "loss" in return_metrics and "loss_loss" in return_metrics:
    del return_metrics["loss_loss"]
  return return_metrics

def test_step(self, inputs):
  input_ids = inputs['input_ids']
  attention_mask = inputs['attention_mask']
  labels = inputs['labels']
  labels_mask = inputs['labels_mask']
  outputs = self(input_ids=input_ids,
                   attention_mask=attention_mask,
                   labels=labels,
                   decoder_attention_mask=labels_mask,
                   training=False
                )
  if not self.loss:
    self.loss_tracker.update_state(y_pred.loss)
    return_metrics = {"loss": self.loss_tracker.result()}
  else:
    return_metrics = {}
  self.compiled_loss(labels, outputs.logits, regularization_losses=self.losses)
  self.compiled_metrics.update_state(labels, outputs.logits)
  for metric in self.metrics:
    result = metric.result()
    if isinstance(result, dict):
        return_metrics.update(result)
    else:
        return_metrics[metric.name] = result
  if "loss" in return_metrics and "loss_loss" in return_metrics:
    del return_metrics["loss_loss"]
  return return_metrics

import functools
t5_model.train_step = functools.partial(train_step, t5_model)
t5_model.test_step = functools.partial(test_step, t5_model)

learning_rate = 0.00005
batch_size = 8
num_epochs = 2
optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=learning_rate)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

t5_model.compile(optimizer=optimizer, 
                 loss=loss_fn,
                 metrics=[tfa.metrics.CohenKappa(num_classes=4,
                                                 weightage='quadratic',
                                                 sparse_labels=True)])

history = t5_model.fit(tokenized_train_data,
             validation_data=tokenized_test_data,
             callbacks=None,
             batch_size=batch_size,
             epochs=num_epochs)
```

### Expected behavior

The model should return results instead of producing a crash.
Looks like the metric does not like your labels. Are you sure it can be used in that case? cc @Rocketknight1 <|||||>Thank you @sgugger for your response. I was hoping it would work, I ran this same metric through a bert model and using its own tokenizers and I had no issues with it. Is there a way to tweak the fit function to insert the labels as is instead of a tensor?<|||||>I believe the issue is caused by your combination of model and metric. `TFT5ForConditionalGeneration` is a model that outputs text, where the distribution over output tokens is conditioned on some input text. Tasks that are suitable for conditional generation models include summarization and translation. 

When using any model that generates text, the number of output classes is equal to the vocabulary size of the model - the model produces a distribution over all possible tokens at each position. However, your metric uses `num_classes=4`. This results in an error because the vocabulary for T5 is thousands of tokens, and so label values can be much higher than 4.

If this code worked with a BERT model, this is because BERT models mostly do not generate text. If you used e.g. `TFBertForSequenceClassification` or `TFBertForTokenClassification`, then the number of classes would be much lower. That is because these models predict categories for each token or for the entire sequence, and the number of categories is set by the `num_labels` argument to the model's `from_pretrained` method.<|||||>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.