TypeError: Cannot convert value <tensorflow.python.keras.losses.CategoricalCrossentropy object ...> to a TensorFlow DType

I want to implement a Word2Vec using negative sampling with pure TensorFlow 2. When I want to compute the gradient I get this error in the last line. I'm struggling to find the problem.

the code is fairly simple:

import tensorflow as tf
import numpy as np

x, y = (('self', 'the'), ('self', 'violent'), ('self', 'any')), (1, 0, 0)
y = tf.convert_to_tensor(y, dtype='float32')

embeding_tensor = tf.keras.layers.Embedding(len(words_lst), embeding_size)
context_tensor = tf.keras.layers.Embedding(len(words_lst), embeding_size)


with tf.GradientTape() as t:    
      middle = embeding_tensor(word2index[x[0][0]])
      neighbor_choices = context_tensor(np.asarray([[word2index[i[1]] for i in x]]))

      scores = tf.tensordot(neighbor_choices, middle, 1)
      prediction = tf.nn.sigmoid(scores)
      loss = tf.keras.losses.CategoricalCrossentropy(y, prediction)

      g_embed, g_context = t.gradient(loss, [middle, neighbor_choices])

TypeError                                 Traceback (most recent call last)
<ipython-input-40-fba4cda17cff> in <module>()
     17       loss = tf.keras.losses.CategoricalCrossentropy(y, prediction)
     18 
---> 19       g_embed, g_context = t.gradient(loss, [middle, neighbor_choices])

2 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py in as_dtype(type_value)
    648 
    649   raise TypeError("Cannot convert value %r to a TensorFlow DType." %
--> 650                   (type_value,))

TypeError: Cannot convert value <tensorflow.python.keras.losses.CategoricalCrossentropy object at 0x7f3ec9be28d0> to a TensorFlow DType.