Implementing a batch dependent loss in Keras

I have an autoencoder set up in Keras. I want to be able to weight the features of the input vector according to a predetermined 'precision' vector. This continuous valued vector has the same length as the input, and each element lies in the range [0, 1], corresponding to the confidence in the corresponding input element, where 1 is completely confident and 0 is no confidence.

I have a precision vector for every example.

I have defined a loss that takes into account this precision vector. Here, reconstructions of low-confidence features are down-weighted.

def MAEpw_wrapper(y_prec):
    def MAEpw(y_true, y_pred):
        return K.mean(K.square(y_prec * (y_pred - y_true)))
    return MAEpw

My issue is that the precision tensor y_prec depends on the batch. I want to be able to update y_prec according to the current batch so that each precision vector is correctly associated with its observation.

I have the done the following:

global y_prec
y_prec = K.variable(P[:32])

Here P is a numpy array containing all precision vectors with the indices corresponding to the examples. I initialize y_prec to have the correct shape for a batch size of 32. I then define the following DataGenerator:

class DataGenerator(Sequence):

    def __init__(self, batch_size, y, shuffle=True):

        self.batch_size = batch_size

        self.y = y

        self.shuffle = shuffle
        self.on_epoch_end()

    def on_epoch_end(self):
        self.indexes = np.arange(len(self.y))

        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __len__(self):
        return int(np.floor(len(self.y) / self.batch_size))

    def __getitem__(self, index):
        indexes = self.indexes[index * self.batch_size: (index+1) * self.batch_size]

        # Set precision vector.
        global y_prec
        new_y_prec = K.variable(P[indexes])
        y_prec = K.update(y_prec, new_y_prec)

        # Get training examples.
        y = self.y[indexes]

        return y, y

Here I am aiming to update y_prec in the same function that generates the batch. This seems to be updating y_prec as expected. I then define my model architecture:

dims = [40, 20, 2]

model2 = Sequential()
model2.add(Dense(dims[0], input_dim=64, activation='relu'))
model2.add(Dense(dims[1], input_dim=dims[0], activation='relu'))
model2.add(Dense(dims[2], input_dim=dims[1], activation='relu', name='bottleneck'))
model2.add(Dense(dims[1], input_dim=dims[2], activation='relu'))
model2.add(Dense(dims[0], input_dim=dims[1], activation='relu'))
model2.add(Dense(64, input_dim=dims[0], activation='linear'))

And finally, I compile and run:

model2.compile(optimizer='adam', loss=MAEpw_wrapper(y_prec))
model2.fit_generator(DataGenerator(32, digits.data), epochs=100)

Where digits.data is a numpy array of observations.

However, this ends up defining separate graphs:

StopIteration: Tensor("Variable:0", shape=(32, 64), dtype=float32_ref) must be from the same graph as Tensor("Variable_4:0", shape=(32, 64), dtype=float32_ref).

I've scoured SO for a solution to my problem but nothing I've found works. Any help on how to do this properly is appreciated.

from Implementing a batch dependent loss in Keras