Joblib can't seem to parallelize this code for saliency maps calculation

Question

def smoothgrad(model, image, actual, noise_std=0.01, num_samples=30):
    
    image = tf.convert_to_tensor(image, dtype=tf.float32)
    
    for _ in range(num_samples):
        
        Smap = []
        total = np.zeros((image.shape[1], image.shape[2]))
        noise = np.random.normal(0, noise_std)
        
        for i in range(image.shape[1]):
            perturbed_row_image = tf.identity(image).numpy()
            perturbed_row_image[:, i, :, :] += noise
            perturbed_row_image = tf.convert_to_tensor(perturbed_row_image)
            
            saliency_map = get_saliency_map_row(model, perturbed_row_image, actual)
            
            total += saliency_map
            
        Smap.append(total)

    smoothgrads = np.mean(Smap, axis=0)

    # Normalize the total saliency map
    min_val = np.min(smoothgrads)
    max_val = np.max(smoothgrads)
    smoothgrads = (smoothgrads - min_val) / (smoothgrads - min_val + 1E-8)

    return smoothgrads

def get_saliency_map_row(model, image, actual):
    
    
    with tf.GradientTape() as tape:
        tape.watch(image)
        predictions = model(image, training=False)
        loss = tf.keras.losses.mean_squared_error(actual, predictions)

    gradient = tape.gradient(loss, image)
    saliency_map = tf.abs(gradient)

    return saliency_map[0, :, :, 0].numpy()




def process_image(i, pbar):
    temp_maps = []
    
    for cnn in CNNs:
        temp_maps.append(smoothgrad(cnn, X_image[i:i+1], y_image[i], num_samples=20))
    
    average_map = np.average(temp_maps, axis=0)
    median_map = np.median(temp_maps, axis=0)
    
    pbar.update()
    
    return (average_map, median_map)

num_jobs = 20


with tqdm(total=len(image_index)) as pbar:
    results = Parallel(n_jobs=num_jobs)(delayed(process_image)(i, pbar) for i in range(len(image_index)))

The provided code implements SmoothGrad saliency maps computation for multiple images in parallel. The smoothgrad function perturbs input images with noise, computes saliency maps for each perturbed image, and averages them to obtain a single SmoothGrad map. The get_saliency_map_row function calculates saliency maps for individual rows of pixels using TensorFlow's GradientTape. The process_image function processes each image in parallel, computing SmoothGrad maps for each model in the list CNNs, and then averaging and computing the median of these maps. Utilizing the joblib library for parallelization, the code distributes workload across multiple CPU cores for efficient computation, providing both average and median SmoothGrad maps for each image.

I expected this code to compute the saliency maps in parallel based on a list of models defined in CNNs. Finally, it should average and compute the median of these maps to obtain a single map for each image or dataset.