Using SSIM loss in TensorFlow returns NaN values

I'm training a network with MRI images and I wanted to use SSIM as loss function. Till now I was using MSE, and everything was working fine. But when I tried to use SSIM (tf.image.ssim), I get a bunch of these warining messages:

 /usr/local/lib/python3.7/dist-packages/matplotlib/image.py:397: UserWarning: Warning: converting a masked element to nan.
  dv = (np.float64(self.norm.vmax) -
/usr/local/lib/python3.7/dist-packages/matplotlib/image.py:398: UserWarning: Warning: converting a masked element to nan.
  np.float64(self.norm.vmin))
/usr/local/lib/python3.7/dist-packages/matplotlib/image.py:405: UserWarning: Warning: converting a masked element to nan.
  a_min = np.float64(newmin)
/usr/local/lib/python3.7/dist-packages/matplotlib/image.py:410: UserWarning: Warning: converting a masked element to nan.
  a_max = np.float64(newmax)
/usr/local/lib/python3.7/dist-packages/matplotlib/colors.py:933: UserWarning: Warning: converting a masked element to nan.
  dtype = np.min_scalar_type(value)
/usr/local/lib/python3.7/dist-packages/numpy/ma/core.py:713: UserWarning: Warning: converting a masked element to nan.
  data = np.array(a, copy=False, subok=subok)

I code is running anyway but no figure is being produced. I am not sure what's happening here or where should I look. I am using tensorflow 2.4.0.

I am attaching a summary of my code my code here:

generator = Generator()  #An u-net defined in tf.keras

gen_learningrate = 5e-4
generator_optimizer = tf.keras.optimizers.Adam(gen_learningrate, beta_1=0.9, beta_2=0.999, epsilon=1e-8)
# Generator loss

def generator_loss(gen_output, target):
    # SSIM loss
    loss = - tf.reduce_mean(tf.image.ssim(target, gen_output, 2)) 
    return loss

@tf.function
def train_step(input_image, target, epoch):
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        gen_output = generator(input_image, training=True)
        
        loss = generator_loss(gen_output, target)

    generator_gradients = gen_tape.gradient(loss, generator.trainable_variables)
    generator_optimizer.apply_gradients(zip(generator_gradients, 
                                                 generator.trainable_variables))

    return loss

def fit(train_ds, epochs, test_ds):
    for input_image, target in train_ds:
            loss = train_step(input_image,target,epoch) 

fit(train_dataset, EPOCHS, test_dataset)

I have explored a little bit and noticed most people using tf.image.ssim() as loss function have used tf.train() from tensorflow 1.0 or model.fit() from tf.keras. I suspect the NaN value returned has something to do with GradientTape() function but I'm not sure how.