Invalid Argument Error: Graph Execution Error on Tensorflow NN

I'm trying to fit the following neural network:

def make_model():
  input =  tf.keras.Input(shape=train_df.shape[1:])
  x = tf.keras.layers.Flatten()(input)
  x = tf.keras.layers.Dense(128, activation='relu')(x)
  x = tf.keras.layers.Dense(64, activation='relu')(x)
  x = tf.keras.layers.Dense(32, activation='relu')(x)
  output = tf.keras.layers.Dense(8, activation='softmax')(x)
  model = tf.keras.models.Model(input,output)
  
  model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
                loss= tf.keras.losses.CategoricalCrossentropy(),
                metrics=[tf.keras.metrics.CategoricalAccuracy(),
                        tf.keras.metrics.FalseNegatives(),
                        tf.keras.metrics.AUC(name='prc', curve='PR')])

  return model
model = make_model()
model.fit(x=train_features, y=train_labels, epochs=2)

where:

model.summary()
print(train_features.shape, train_labels.shape)

outputs the following:

Model: "model_8"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_9 (InputLayer)        [(None, 17)]              0         
                                                                 
 flatten_8 (Flatten)         (None, 17)                0         
                                                                 
 dense_32 (Dense)            (None, 128)               2304      
                                                                 
 dense_33 (Dense)            (None, 64)                8256      
                                                                 
 dense_34 (Dense)            (None, 32)                2080      
                                                                 
 dense_35 (Dense)            (None, 8)                 264       
                                                                 
=================================================================
Total params: 12,904
Trainable params: 12,904
Non-trainable params: 0
_________________________________________________________________
Train Features Shape: (64140, 17) 
Train Labels Shape: (64140, 8)

However, it keeps getting this error mid-epoch:

    Epoch 1/2
 444/2005 [=====>........................] - ETA: 1s - loss: 1.1139 - categorical_accuracy: 0.4904 - false_negatives_6: 10143.0000 - prc: 0.5232
Output exceeds the size limit. Open the full output data in a text editor
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
c:\Users\nrtc\OneDrive\Documentos\AI Summer School\Competition\Cópia_de_imbalanced_data.ipynb Cell 34' in <module>
----> 1 model.fit(x=train_features, y=train_labels, epochs=2)

File c:\Python39\lib\site-packages\keras\utils\traceback_utils.py:67, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     65 except Exception as e:  # pylint: disable=broad-except
     66   filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67   raise e.with_traceback(filtered_tb) from None
     68 finally:
     69   del filtered_tb

File ~\AppData\Roaming\Python\Python39\site-packages\tensorflow\python\eager\execute.py:54, 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:
     57   if name is not None:

InvalidArgumentError: Graph execution error:

Detected at node 'assert_greater_equal/Assert/AssertGuard/Assert' defined at (most recent call last):
    File "c:\Python39\lib\runpy.py", line 197, in _run_module_as_main
      return _run_code(code, main_globals, None,
...
    File "c:\Python39\lib\site-packages\keras\utils\metrics_utils.py", line 602, in update_confusion_matrix_variables
      tf.debugging.assert_greater_equal(
Node: 'assert_greater_equal/Assert/AssertGuard/Assert'
assertion failed: [predictions must be >= 0] [Condition x >= y did not hold element-wise:] [x (model_6/dense_27/Softmax:0) = ] [[0.101199746 0.358387947 0.118633337...]...] [y (Cast_4/x:0) = ] [0]
     [[{{node assert_greater_equal/Assert/AssertGuard/Assert}}]] [Op:__inference_train_function_10341]

Any idea what might be the error? I've seen other threads for the same error (*), but I do think that the last layer is correctly set up for 8 output labels.

*Other stack overflow threads for a similar problem

https://stackoverflow.com/questions/62606345/tensorflow-2-2-0-error-predictions-must-be-0-condition-x-y-did-not-hold

https://stackoverflow.com/questions/71153492/invalid-argument-error-graph-execution-error