KerasMetricCallback expecting dictionary but receiving numpy array
### System Info

Running in Google Colab with `!pip install transformers evaluate` as the first cell. The results of `transformers-cli env` are:

- `transformers` version: 4.26.1
- Platform: Linux-5.10.147+-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.12.1
- PyTorch version (GPU?): 1.13.1+cu116 (False)
- Tensorflow version (GPU?): 2.11.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

### 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

[View the Google Colab here](https://colab.research.google.com/drive/1Pgc1jkZZbMmOF4O8Tz4Nz81V5FqkqCr-?usp=sharing)

Code:

```
!pip install transformers evaluate

import tensorflow as tf
import evaluate
from numpy import argmax as np_argmax
from transformers import create_optimizer
from transformers.keras_callbacks import KerasMetricCallback

tf.debugging.disable_traceback_filtering()

train_texts = ["This is class 0", "I am a class 1 sentence", "Class 2", "Also class 2"]
train_labels = [0, 1, 2, 2]

test_texts = ["A class 1 example", "Testing class 0"]
test_labels = [1, 0]

num_classes = 3

batch_size = 16

def create_dataset(texts, labels):
  dataset = tf.data.Dataset.from_tensor_slices((texts, labels))
  return dataset.shuffle(10000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

train_dataset = create_dataset(train_texts, train_labels)
test_dataset = create_dataset(test_texts, test_labels)

encoder = tf.keras.layers.TextVectorization()
encoder.adapt(train_dataset.map(lambda text, label: text))

model = tf.keras.Sequential([
    encoder,
    tf.keras.layers.Embedding(
        input_dim=len(encoder.get_vocabulary()),
        output_dim=64,
        # Use masking to handle the variable sequence lengths
        mask_zero=True),
    tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(num_classes, activation='sigmoid')
])

num_epochs = 5
batches_per_epoch = len(train_texts) // batch_size
total_train_steps = int(batches_per_epoch * num_epochs)
optimizer, schedule = create_optimizer(init_lr=2e-5, num_warmup_steps=0, num_train_steps=total_train_steps)

model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                  optimizer=optimizer)

accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np_argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=test_dataset)
callbacks = [metric_callback]

model.fit(x=train_dataset, validation_data=test_dataset, epochs=num_epochs, callbacks=callbacks)
```

Full error stack trace:
```
Epoch 1/5
1/1 [==============================] - ETA: 0s - loss: 1.0996
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-17-e77f61379ec7>](https://localhost:8080/#) in <module>
----> 1 model.fit(x=train_dataset, validation_data=test_dataset, epochs=num_epochs, callbacks=callbacks)

3 frames
[/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py](https://localhost:8080/#) in error_handler(*args, **kwargs)
     59     def error_handler(*args, **kwargs):
     60         if not tf.debugging.is_traceback_filtering_enabled():
---> 61             return fn(*args, **kwargs)
     62 
     63         filtered_tb = None

[/usr/local/lib/python3.8/dist-packages/keras/engine/training.py](https://localhost:8080/#) in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
   1710                     epoch_logs.update(val_logs)
   1711 
-> 1712                 callbacks.on_epoch_end(epoch, epoch_logs)
   1713                 training_logs = epoch_logs
   1714                 if self.stop_training:

[/usr/local/lib/python3.8/dist-packages/keras/callbacks.py](https://localhost:8080/#) in on_epoch_end(self, epoch, logs)
    452         logs = self._process_logs(logs)
    453         for callback in self.callbacks:
--> 454             callback.on_epoch_end(epoch, logs)
    455 
    456     def on_train_batch_begin(self, batch, logs=None):

[/usr/local/lib/python3.8/dist-packages/transformers/keras_callbacks.py](https://localhost:8080/#) in on_epoch_end(self, epoch, logs)
    236                     predictions = {key: predictions[key] for key in self.output_cols}
    237                 else:
--> 238                     predictions = {key: val for key, val in predictions.items() if key not in ignore_keys + ["loss"]}
    239             prediction_list.append(predictions)
    240             if not self.use_keras_label:

AttributeError: 'numpy.ndarray' object has no attribute 'items'
```

### Expected behavior

The code is adapted from a [HuggingFace text classification tutorial](https://huggingface.co/docs/transformers/tasks/sequence_classification#text-classification) and a [TensorFlow text classification with an RNN tutorial](https://www.tensorflow.org/text/tutorials/text_classification_rnn). The optimizer, metrics and callbacks are from the HuggingFace tutorial. The encoder and model are from the TensorFlow tutorial.

The error given is `AttributeError: 'numpy.ndarray' object has no attribute 'items'` and occurs from [line 237 of keras_callbacks.py](https://github.com/huggingface/transformers/blob/main/src/transformers/keras_callbacks.py#L237). The code around this seems to expect to be dealing with a dictionary or a subclass of a dictionary as a result of the `predict_on_batch` function called on the model. However, [the documentation for the TensorFlow Model class](https://www.tensorflow.org/api_docs/python/tf/keras/Model#predict_on_batch), [the documentation for the TensorFlow Sequential Class (which subclasses Model) ](https://www.tensorflow.org/api_docs/python/tf/keras/Sequential#predict_on_batch) and [the source code for the `predict_on_batch` method](https://github.com/keras-team/keras/blob/v2.11.0/keras/engine/training.py#L2547-L2572) show that it returns a numpy array.

I would expect this code not to error, and the callback to successfully call the metrics function with the expected predictions.
cc @Rocketknight1 <|||||>Hi @leadbetterben, the problem arises because the metric callback was intended for use with `transformers` models, which generally return dicts or tuples of outputs rather than just a single array. This was an oversight on my part - I'll see if I can push a fix!<|||||>@leadbetterben I've created [a PR](https://github.com/huggingface/transformers/pull/21727) to resolve this issue. Can you try it out? To use the PR branch, replace the first block in your Colab notebook with this:
```
!pip install git+https://github.com/huggingface/transformers.git@metric_callback_fix
!pip install evaluate
```<|||||>@Rocketknight1 I've just tested it out and it seems to work fine. Thank you for that 👍 <|||||>@leadbetterben The PR has now been merged. You can use it by installing `transformers` from `main` with `!pip install git+https://github.com/huggingface/transformers.git`. It'll also be included in our next release, after which you can go back to just using `pip install transformers`.

Thanks again for the bug report!<|||||>Thank you @Rocketknight1 !