None gradient using custom loss function that uses MSSIM

TensorFlow General Discussion
, ,
1

Hello, I am a novice with tensorflow, I am trying create a custom loss function that uses MSSIM. I am trying to calculate the MSSIM of a ground truth with the gradcam++ sailancy map of the input image.
Bellow is some of the code snipits.

Map calculation:
def calc_maps(model, HE_images, HE_labels):
HE_labels = list(tf.argmax(HE_labels, axis=1))
replace2linear = ReplaceToLinear()
score = CategoricalScore(HE_labels)
gradcam = GradcamPlusPlus(model, model_modifier=replace2linear, clone=True)
cam = tf.convert_to_tensor(gradcam(score,HE_images,penultimate_layer=-1))

return cam

A trail I did:
HE, IHC = train_data_loader[0]
LABELS = tf.convert_to_tensor(HE[-1])
HE = tf.convert_to_tensor(HE[:-1][0])
IHC = tf.convert_to_tensor(IHC[:-1][0])

MAPS = 1 - calc_maps(model, HE, LABELS)

with tf.GradientTape() as tape:
predect = model(HE) #forward

im1 =tf.image.convert_image_dtype(IHC, tf.float32)
im2 =tf.image.convert_image_dtype(MAPS, tf.float32)

im2 = tf.expand_dims(im2, axis=-1)

ssim = tf.image.ssim_multiscale(im1, im2, 1.0)

gradients = tape.gradient(ssim, model.trainable_variables)

print(gradients)