TensorFlow 2 custom loss: “No gradients provided for any variable” error

Question

I trying to run my module with custom loss function Here is my loss function:

def softargmax(x, beta=1e10):
    x_range = tf.range(x.shape.as_list()[-1], dtype=x.