I am trying to write a custom loss function to a Keras model. This loss function takes in the prediction and suppresses all the predictions except the highest one to zero. Something like this:
def loss_function(y_true, y_pred):
s=tf.shape(y_pred)
loop=tf.unstack(y_pred)
th=[tf.math.argmax(y_pred[i]) for i in loop]
output=tf.zeros([len(loop),len(anchor_list)])
y_pred_in=tf.tensor_scatter_nd_update(output,[[i,th[i]] for i in loop],[y_pred[i,th[i]] for i in loop])
cce = tf.keras.losses.CategoricalCrossentropy()
loss=cce(y_true,y_pred)
return loss
I realized that while doing so I have to take care of batch size as well. For this, I want to create lists and arrays according to the shape of the prediction tensor. I found out that we cannot iterate over a tensor. So, I unstacked the tensor. However, get the following error:
shape=tf.unstack(y_pred)
ValueError: Cannot infer argument `num` from the shape (None, None)
How do I solve this problem?