Element not part of the graph when using VGG for data generation and loss calculation

I have a VGG19 encoder which takes an input image y of (256,256,3) and returns a tensor of dimension (32,32, 512) from conv-4-1 layer of vgg. I need to turn it into a numpy array to apply some transformations and reconstruct the image using my decoder.

In short, I'm trying to train the decoder network like so:

x = vgg_encoder(y)  # generate features from image y
x = do_extra_transformation(x) # for example, reshape and apply K means to shift features towards their cluster centres
y_pred = decoder(x) # try to reconstruct the image y from features
loss = calculate_loss(y, y_pred) # calculate reconstruction loss using VGG loss

However, when I run the code, I get the error: ValueError: Tensor Tensor("block4_conv1/Relu:0", shape=(?, 32, 32, 512), dtype=float32) is not an element of this graph.

I'm assuming the error comes from tensorflow diconnecting the graph after I call predict on VGG to generate features. I don't see why this is a problem, since it technically only used for data generation, and isn't part of the computation graph for training!

---

Full code, you can run with python example.py below

import tensorflow as tf
import numpy as np 
from tensorflow.keras.applications import VGG19
from tensorflow.keras.layers import Input, UpSampling2D, Conv2D
from tensorflow.keras.models import Model
import tensorflow.keras.backend as K
from tensorflow.keras.optimizers import Adam

class CustomModel:

    def __init__(self, im_h, im_w, im_c):
        self.im_shape = (im_h, im_w, im_c)
        self.vgg_features_shape = (None, None, 512)
        self.vgg_loss_model = self.build_vgg_loss()
        self.kernel_size = (3,3)
        self.decoder = self.build_decoder()


    def build_vgg_loss(self):
        vgg = VGG19(weights="imagenet", include_top=False, input_shape=self.im_shape)
        vgg.outputs = vgg.get_layer('block4_conv1').output
        model = Model(inputs=vgg.inputs, outputs=vgg.outputs)
        model.trainable = False

        return model

    def build_decoder(self):
        """
        Mirrors the VGG network with max-pooling layers replaces by UpScaling Layers
        """

        i = Input((None, None, 512))
        x = Conv2D(filters=512, kernel_size=self.kernel_size, padding='same')(i)

        x = UpSampling2D()(x)
        for _ in range(4):
            x = Conv2D(filters=256, kernel_size=self.kernel_size, padding='same')(x)

        x = UpSampling2D()(x)
        for _ in range(2):
            x = Conv2D(filters=128, kernel_size=self.kernel_size, padding='same')(x)

        x = UpSampling2D()(x)
        for _ in range(2):
            x = Conv2D(filters=64, kernel_size=self.kernel_size, padding='same')(x)

        x = Conv2D(filters=3, kernel_size=self.kernel_size, padding='same')(x)

        model = Model(inputs=i, outputs=x)

        return model


    def get_loss(self, y_pred, y):

        vgg_model = self.vgg_loss_model

        def content_loss(y_pred, y):

            dif = vgg_model(y) - vgg_model(y_pred)
            sq = K.square(dif)
            s = K.sum(sq, axis=-1)
            sqrt = K.sqrt(s)
            loss = K.sum(sqrt)

            return loss

        return content_loss(y_pred, y)


class DataLoader:

    def __init__(self, vgg):
        self.vgg = vgg

    def gen(self):
        while True:
            y = np.random.randn(256, 256,3)
            x = self.vgg.predict(np.expand_dims(y, 0)).reshape((32,32,512)) # if this is turned into a np.array, everything works as expected
            yield x, np.random.randn(256, 256,3) 


model = CustomModel(256,256,3)

# dl = DataLoader(datapath='./trainer/data/', mst=mst)

output_types=(
    tf.float32,
    tf.float32
 )
output_shapes=(
    tf.TensorShape([None, None, None]),
    tf.TensorShape([None, None, None])
)

ds = tf.data.Dataset.from_generator(DataLoader(model.vgg_loss_model).gen,
                                   output_types=output_types, 
                                   output_shapes=output_shapes)

ds = ds.repeat().batch(1)
iterator = ds.make_one_shot_iterator()
x, y = iterator.get_next()
y_pred = model.decoder(x)


loss = model.get_loss(y_pred, y)
opt = tf.train.AdamOptimizer(0.01)
train_opt = opt.minimize(loss)

init_op = tf.global_variables_initializer()

with tf.Session() as sess:

    sess.run(init_op)
    opt = tf.train.GradientDescentOptimizer(0.01)
    for i in range(5):
        sess.run(train_opt)

from Element not part of the graph when using VGG for data generation and loss calculation