Custom loss function with both min(y, p) and max(y,p)
A loss function calculates the error over all the data presented to it. For neural networks, that is an average over the mini-batch. Your code might only handle scalars (hence working when batch=1), it should handle vectors.
The K looks like you are using tf.Keras. The documentation gives examples of custom loss functions. It helps to subclasses tf.keras.losses.Loss which expects a call() method that contains the logic for loss calculation using y_true, y_pred.
class CustomLoss(Loss):

  def call(self, y_true, y_pred):
    pass