How can I discretize the output of a deep learning model?

Question

I am trying to create a supervised trained model that outputs discrete values ranging from 1 to 50. To do this, I decided to solve the problem as a multiclassification problem. My approach is to use a dense layer of 50 neurons with Softmax activation function at the end of the model, and then also a Lambda function to extract the value from the array.

model_3=tf.keras.Sequential([
tf.keras.layers.Conv1D(128,1,padding="same",activation="relu",input_shape= (WIN_LEN,FEATURES)),
tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(50,activation="softmax")),
tf.keras.layers.TimeDistributed(tf.keras.layers.Lambda(lambda x: x.argmax()+1))])

However, I am receiving an error that says:

Call arguments received by layer "lambda_5" "f"(type Lambda):
  • inputs=tf.Tensor(shape=(None, 50), dtype=float32)
  • mask=None
  • training=None

Does anyone have an idea why it's not working? This question is a continuation of another question.