Can't add other metrics than accuracy on ViT model

I am a begginer in machine learning and i am trying to train a ViT model to categorical classes with my own dataset. I am following this code: https://keras.io/examples/vision/image_classification_with_vision_transformer/

It’s working fine when I use the accuracy metric, but I want to use recall and precision as well, but I keep getting this error when I add them:

/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52   try:
     53     ctx.ensure_initialized()
---> 54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:

InvalidArgumentError: Graph execution error:

Detected at node 'assert_greater_equal/Assert/AssertGuard/Assert' defined at (most recent call last):
    File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
      exec(code, run_globals)
    File "/usr/local/lib/python3.8/dist-packages/ipykernel_launcher.py", line 16, in <module>
      app.launch_new_instance()
    File "/usr/local/lib/python3.8/dist-packages/traitlets/config/application.py", line 992, in launch_instance
      app.start()
    File "/usr/local/lib/python3.8/dist-packages/ipykernel/kernelapp.py", line 612, in start
      self.io_loop.start()
    File "/usr/local/lib/python3.8/dist-packages/tornado/platform/asyncio.py", line 149, in start
      self.asyncio_loop.run_forever()
...

(1) INVALID_ARGUMENT:  assertion failed: [predictions must be >= 0] [Condition x >= y did not hold element-wise:] [x (model_3/dense_79/BiasAdd:0) = ] [[251.636795 -233.491394 322.750397...]...] [y (Cast_4/x:0) = ] [0]
     [[{{node assert_greater_equal/Assert/AssertGuard/Assert}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_348766]

I also encoded my y_train to one_hot so I could use the Categorical_Crossentropy loss instead of SparseCategorical_Crossentropy loss

Here is the shape of the arrays now:

x_train shape: (6179, 336, 336, 3) – y_train shape: (6179, 9) x_test shape: (2060, 336, 336, 3) – y_test shape: (2060, 9)

I just changed a few things on the compile from the original code:

optimizer = tfa.optimizers.AdamW(
    learning_rate=learning_rate, weight_decay=weight_decay
)

model.compile(
        optimizer=optimizer,
        loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
        metrics = [ 'accuracy', tf.keras.metrics.Recall(name = 'recall')]
)

checkpoint_filepath = "/content/drive/MyDrive/DATASET/checkpoints/checkpoint.h5"
checkpoint_callback = keras.callbacks.ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_accuracy",
        save_best_only=True,
        save_weights_only=True,
    )

r = model.fit(
        x=x_train,
        y=y_train,
        batch_size=batch_size,
        epochs=num_epochs,
        validation_split=0.25,
)

model.load_weights(checkpoint_filepath)
_, accuracy = model.evaluate(x_test, y_test)
print(f"Test accuracy: {round(accuracy * 100, 2)}%")

Im using colab