So I’d like to create a custom Tensorflow.Keras layer with purpose of returning label name as a text.
The previous layer is an argmax that returns integer representing the index of predicted class and I want to pass it as input to the next layer which returns the real name of the class as the final output of the model.

This is my custom layer

class LabelLayer(tf.keras.layers.Layer):
  def __init__(self, labels, dtype='int16'):
    super(LabelLayer, self).__init__()
    self.labels = labels

  def call(self, inputs):
    return self.labels[inputs]

And this is how I build the model, on top of an existing model

new_model = tf.keras.Sequential([
  model,
  tf.keras.layers.Lambda(lambda x: tf.cast( tf.argmax(x[0]), dtype='int16'), name='predicted_label'),
  LabelLayer(class_labels) # TODO
])

Right now it throws an error:
“TypeError: list indices must be integers or slices, not Tensor”
I tried to using some suggested solutions such as tf.slice or invoking .numpy() method but it doesn’t seem to work.
What am I doing wrong? Is there a more correct way to achieve what I want?