I have CNN model build in tensorflow which outputs grid with shape (batch_size,H,W,1).
i perform sigmoid activation function on each cell in the grid.
my y_true is tensor with [ColIndex,RowIndex,Outcome], meaning for every training instance i have exactly 1 cell with an outcome value.

i want to calculate the loss when training only for the 1 cell in the y_true, meaning i need to index the grid output to this 1 cell and calculate logloss, and use this 1 cell loss when training.

i have looked for solutions but isn’t found one, i currently have managed to create a loss which the model does train but the loss are almost always 0 and i get weird results.

my loss function:

def custom_target_location_loss(y_true, y_pred):

    # Extract column indices, row indices, and values
    col_indices = tf.cast(y_true[:,0], tf.int32)  # Column indices
    row_indices = tf.cast(y_true[:,1], tf.int32)  # Row indices
    values = tf.cast(y_true[:, 2], tf.float32)     # Binary labels (0 or 1)
    
    batch_size = tf.shape(y_pred)[0]
    batch_indices = tf.range(batch_size)
    channel_indices = tf.zeros_like(batch_indices, dtype=tf.int32)
    indices = tf.stack([batch_indices,row_indices, col_indices,channel_indices], axis=1)
    
    
    y_pred_cell = tf.gather_nd(y_pred, indices)
    
    loss = tf.keras.losses.binary_crossentropy(values, y_pred_cell)

    # Return the mean loss across the batch
    return tf.reduce_mean(loss)

am i missing something here?