Custom loss function implementation issue in keras

I am implementing a custom loss function in keras. The output of the model is 10 dimensional softmax layer. To calculate loss: first I need to find the index of y firing 1 and then subtract that value with true value. I'm doing the following:

from keras import backend as K

def diff_loss(y_true,y_pred):

    # find the indices of neuron firing 1
    true_ind=K.tf.argmax(y_true,axis=0)
    pred_ind=K.tf.argmax(y_pred,axis=0)

    # cast it to float32
    x=K.tf.cast(true_ind,K.tf.float32)
    y=K.tf.cast(pred_ind,K.tf.float32)

    return K.abs(x-y)

but it gives error "raise ValueError("None values not supported.") ValueError: None values not supported." What's the problem here?

This happens because your function is not differentiable. It's made of constants.

There is simply no solution for this if you want argmax as result.

An approach to test

Since you're using "softmax", that means that only one class is correct (you don't have two classes at the same time).

And since you want index differences, maybe you could work with a single continuous result (continuous values are differentiable).

Work with only one output ranging from -0.5 to 9.5, and take the classes by rounding the result.

That way, you can have the last layer with only one unit:

lastLayer = Dense(1,activation = 'sigmoid', ....) #or another kind if it's not dense

And change the range with a lambda layer:

lambdaLayer = Lambda(lambda x: 10*x - 0.5)

Now your loss can be a simple 'mae' (mean absolute error).

The downside of this attempt is that the 'sigmoid' activation is not evenly distributed between the classes. Some classes will be more probable than others. But since it's important to have a limit, it seems at first the best idea.

This will only work if you classes follow a logical increasing sequence. (I guess they do, otherwise you'd not be trying that kind of loss, right?)