replace elements of prediction in custom loss function for tensor flow

Question

I am using a classification model, but would like to write my custom loss function which considers the value as 1 for two of the three categories only if the softmax value is gr