How can I implement a custom loss function that takes into account multiple predictions of the network?

I am currently implementing a CNN with a custom error function. The problem I am trying to solve is physics-based, so I can calculate the maximal achievable precision, or to put it another way, I know the best possible (i.e. minimal) standard deviation I can achieve. Those best possible precisions are calculated during the generation of the training data using the Cramer-Rao-lower bound (CRLB).

Right now, my error function looks something like this (in Keras):

def customLoss(yTrue, yPred):
    STD = yTrue[:, 10:20]
    yTrue = yTrue[:, 0:10]
    dev = K.mean(K.abs(K.abs(yTrue - yPred) - STD))
    return dev

In this case, I have 10 parameters, so I want to estimate with 10 CRLB's. I put the CRLB's in the target vector just to be able to handle the in the error function.

To my question. This method works, but it is not what I want. The problem is that the error is calculated considering a single prediction of the network, but to be correct the network would have to predict the same dataset/batch multiple times. By doing that I would be able to see the standard deviation of the prediction and use that to calculate the error (I'm using a Bayesian CNN).

Has someone an idea how to implement such a function in Keras or Tensorflow (I would also not mind switching to PyTorch)?