Implementation of statistical constraints inside GAN loss function (Keras)

I'm trying to implement a keras version of the Generative Adversarial Network (GAN) discussed in

Wu, Jin-Long, et al. "Enforcing statistical constraints in generative adversarial networks for modeling chaotic dynamical systems." Journal of Computational Physics 406 (2020): 109209.

I already have a working keras implementation of WGANGP and I'm using it to reproduce particles velocities inside a turbulent flow (1d signals, only x-component). The general structure of my model is this (self.gen and self.critic are keras sequential models for the generator and critic, respectively):

        #-------------------------------
        # Construct Computational Graph
        #       for the Critic
        #-------------------------------

        # Freeze generator's layers while training critic
        self.gen.trainable = False

        # Image input (real sample)
        real_img = Input(shape=(self.sig_len,self.channels))

        # Noise input
        z_disc = Input(shape=(self.noise_dim,))
        # Generate image based of noise (fake sample)
        fake_img = self.gen(z_disc)

        # Discriminator determines validity of the real and fake images
        fake = self.critic(fake_img)
        valid = self.critic(real_img)

        # Construct weighted average between real and fake images
        interpolated_img = RandomWeightedAverage(self.batch_size)([real_img, fake_img])
        # Determine validity of weighted sample
        validity_interpolated = self.critic(interpolated_img)

        # Use Python partial to provide loss function with additional
        # 'averaged_samples' argument
        partial_gp_loss = partial(self.gradient_penalty_loss,
                                  averaged_samples=interpolated_img)
        partial_gp_loss.__name__ = 'gradient_penalty' # Keras requires function names

        self.critic_model = Model(inputs=[real_img, z_disc],
                                  outputs=[valid, fake, validity_interpolated])
        self.critic_model.compile(loss=[self.wasserstein_loss,
                                        self.wasserstein_loss,
                                        partial_gp_loss],
                                  optimizer=critic_optimizer,
                                  loss_weights=[1, 1, 10])

        #-------------------------------
        # Construct Computational Graph
        #         for Generator
        #-------------------------------

        # For the generator we freeze the critic's layers
        self.critic.trainable = False
        self.gen.trainable = True

        # Sampled noise for input to generator
        z_gen = Input(shape=(self.noise_dim,))
        # Generate images based of noise
        img = self.gen(z_gen)
        # Discriminator determines validity
        valid = self.critic(img)

        # Defines generator model
        self.gen_model = Model(inputs=[z_gen], outputs=[valid])
        self.gen_model.compile(loss=self.wasserstein_loss, optimizer=gen_optimizer)

The model is built following this github repository.

I need to add a second loss function to gen_model that computes the covariance matrix of the generated signals (which are 2000 points long) inside the batch and calculates the Frobenius norm of the difference between this computed matrix and the theoretical one, as explained in the image. statistical constraint loss function

I wrote the function using keras backend that computes the covariance matrix and returns the frobenius norm:

    def K_invariance(self,y_true,y_pred):
        
        def K_cov2(x):
            x = K.squeeze(x, axis=-1)
            x = x - K.expand_dims(K.mean(x, axis=0), 0)
            N = K.cast(K.shape(x)[0], dtype="float32")
            return K.dot(K.transpose(x), x) / N

        def K_norm(x):
            return K.sqrt(K.sum(K.square(x)))
        
        cov_theoretical = K.constant(np.load("../data/cij.npy"),shape=(2000,2000))
        cov_computed = K_cov2(x)

        return K_norm(cov_theoretical-cov_computed)

I tried to add this as loss_function to gen_model like this:

        #-------------------------------
        # Construct Computational Graph
        #         for Generator
        #-------------------------------

        # For the generator we freeze the critic's layers
        self.critic.trainable = False
        self.gen.trainable = True

        # Sampled noise for input to generator
        z_gen = Input(shape=(self.noise_dim,))
        # Generate images based of noise
        img = self.gen(z_gen)
        # Discriminator determines validity
        valid = self.critic(img)

        # Defines generator model
        self.gen_model = Model(inputs=[z_gen], **outputs=[valid,img]**)
        self.gen_model.compile(**loss=[self.wasserstein_loss,self.K_invariance],
                               loss_weights=[1,0.001]**, optimizer=gen_optimizer)

The Problem is that during training, the second loss function is computed only the first generator iteration and then remains constant. Is something wrong? That the loss actually remains constant during training is almost impossible because I can see the result of the model during training and they go from complete random noise to beautiful smooth trajectories...