How to use GradCAM for multichannel 1D CNN models

Question

I have some multi-channel time series data which I want to use as input to a 1D Convolutional Neural Network classifier. Furthermore, I want to test the model and provide an activation map for these test data.

I have already implemented a solution that works for single-channel data, but I can't figure out how to make the changes that are necessary in order to get an activation map for multi-channel data.

My single-channel model looks like this (notice that the input shape is (300,1) -> signal length = 300 and number of channels = 1):

def FCN():
    inputlayer = tf.keras.layers.Input(shape=(300,1)) 

    conv1 = tf.keras.layers.Conv1D(filters=128, kernel_size=8,input_shape=(250,12), padding='same')(inputlayer)
    conv1 = tf.keras.layers.BatchNormalization()(conv1)
    conv1 = tf.keras.layers.Activation(activation='relu')(conv1)

    conv2 = tf.keras.layers.Conv1D(filters=256, kernel_size=5, padding='same')(conv1)
    conv2 = tf.keras.layers.BatchNormalization()(conv2)
    conv2 = tf.keras.layers.Activation('relu')(conv2)

    conv3 = tf.keras.layers.Conv1D(128, kernel_size=3,padding='same', name = "last_conv")(conv2)
    conv3 = tf.keras.layers.BatchNormalization()(conv3)
    conv3 = tf.keras.layers.Activation('relu')(conv3)

    gap_layer = tf.keras.layers.GlobalAveragePooling1D()(conv3)


    outputlayer = tf.keras.layers.Dense(1, activation='sigmoid')(gap_layer)

    model = tf.keras.Model(inputs=inputlayer, outputs=outputlayer)

    model.compile(loss=tf.keras.losses.BinaryCrossentropy(), optimizer=tf.keras.optimizers.Adam(learning_rate=0.01), metrics=[tf.keras.metrics.BinaryAccuracy()])

    return model

My grad_cam function for the single-channel model looks like this:

def grad_cam(layer_name, data):
    grad_model = tf.keras.models.Model(
    [model.inputs], [model.get_layer(layer_name).output, model.output]
)
    last_conv_layer_output, preds = grad_model(data)

    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model(data)
        pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]
    
    grads = tape.gradient(class_channel, last_conv_layer_output)

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

    last_conv_layer_output = last_conv_layer_output[0]

    heatmap = last_conv_layer_output * pooled_grads
    heatmap = tf.reduce_mean(heatmap, axis=(1))
    heatmap = np.expand_dims(heatmap,0)
    return heatmap

This is how I first define the model, train it and then use my grad_cam function to provide me an activation map for all data in X_test

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

model = FCN()
model.fit(...)
...
layer_name = "<last conv layer name>"
for i in X_test:
    data = np.expand_dims(i,0)
    heatmap = grad_cam(layer_name,data)

    plt.figure(figsize=(30,4))
    plt.imshow(np.expand_dims(heatmap,axis=2),cmap='Reds', aspect="auto", interpolation='nearest',extent=[250,i.min(),i.max()], alpha=0.5)
    plt.plot(i,'k')
    plt.colorbar()
    plt.show()

But how can all this be achieved for multi-channel data?