Training Heatmap Model - Softmax4D?

[samhains]

When i turn on heatmap = True when training the image classifier I get this error:

fitting model for the dataset
Traceback (most recent call last):
File "image_classification.py", line 122, in
train_network()
File "image_classification.py", line 118, in train_network
shuffle='batch', callbacks=callbacks)
File "/home/sam/anaconda3/envs/barnes/lib/python3.5/site-packages/keras/models.py", line 627, in fit
sample_weight=sample_weight)
File "/home/sam/anaconda3/envs/barnes/lib/python3.5/site-packages/keras/engine/training.py", line 1052, in fit
batch_size=batch_size)
File "/home/sam/anaconda3/envs/barnes/lib/python3.5/site-packages/keras/engine/training.py", line 983, in _standardize_user_data
exception_prefix='model target')
File "/home/sam/anaconda3/envs/barnes/lib/python3.5/site-packages/keras/engine/training.py", line 100, in standardize_input_data
str(array.shape))
Exception: Error when checking model target: expected softmax to have 4 dimensions, but got array with shape (24, 3)

I changed this line to reflect the 3 classes that I am testing this model on:
https://github.com/yardstick17/ConnectingDots/blob/master/neural_networks/convolutional_neural_network/convnets-keras/convnetskeras/convnets.py#L155

However I notice it is the Softmax4D layer causing the problem. I have had a look into the custom layers code to see what Softmax4D is doing, but I am a little unclear about what is going on in there. It would be great if you could clarify what the Softmax4D layer is doing and what I need to do to make it happy!

Sam

[samhains]

i fixed this by adapting code from https://github.com/jacobgil/keras-cam

def global_average_pooling(x):
    return K.mean(x, axis = (2, 3))

def global_average_pooling_shape(input_shape):
    return input_shape[0:2]

    if heatmap:
        model.add(Lambda(global_average_pooling, 
                  output_shape=global_average_pooling_shape))
        model.add(Dense(num_of_output_classes, activation = 'softmax', init='uniform'))