displaying Grad-CAM Heatmap with ResNet50 mode

Question

This is the model summary: model.summary() Last convolution layer is "conv5_block3_out" inside the ResNet layer

Code:

def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):

      grad_model = tf.keras.models.Model(
          [model.get_layer('resnet50').get_layer('input_1').input], [model.get_layer(last_conv_layer_name).get_layer("conv5_block3_out").output,model.get_layer(last_conv_layer_name).output]
      )


      with tf.GradientTape() as tape:
          last_conv_layer_output, preds = grad_model(img_array)
          print("preds:",preds[0])
          if pred_index is None:
              pred_index = tf.argmax(preds[0])
          print("pred_index:",pred_index)
          class_channel = preds[:, pred_index]
          print("class_channel:",class_channel)


      grads = tape.gradient(class_channel, last_conv_layer_output)


      pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

      last_conv_layer_output = last_conv_layer_output[0]
      heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
      heatmap = tf.squeeze(heatmap)

      # For visualization purpose, we will also normalize the heatmap between 0 & 1
      heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
      return heatmap.numpy()
  img=X_Train[4].copy()
  # Convert the image to array of type float32
  img = np.asarray(img, dtype = np.float32)
  img=cv2.resize(img.copy(), (IMAGE_SIZE,IMAGE_SIZE), interpolation=cv2.INTER_AREA)

  img=img.reshape(-1, IMAGE_SIZE, IMAGE_SIZE, 3)
  img_scaled = img / 255


  # Remove last layer's softmax
  model.layers[-1].activation = None


  # # Generate class activation heatmap
  heatmap = make_gradcam_heatmap(img_scaled, model, 'resnet50')**#This line gives an error** "InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [] != values[1].shape = [2,2048] [Op:Pack]"

Output for reference Output

Can somebody correct the grad-cam code, I am stuck at it for so long