Cannot use keras_spiking.ModelEnergy on model that have a concatenation layer, such as DenseNet121

[JavierZazueta]

I am trying to get some energy values from a DenseNet architecture and when I run the code below I get a valueError (as shown below). I tried narrowing down where the issue occurs and found that it first happens at a layers.Concatenate call in the keras implementation. Has anyone seen a similar issue to this? When I use keras’ vgg16 implementation it works, which doesn’t have a Concatenate.

from tensorflow.keras.applications import DenseNet121
import keras_spiking
import numpy as np

denseNet_model = DenseNet121(
weights=None, include_top=True, input_shape=(224, 224, 3)
)

energy = keras_spiking.ModelEnergy(denseNet_model, example_data=np.ones((32, 224, 224, 3)) * 3)

ValueError: could not broadcast input array from shape (32,56,56,64) into shape (32,56,56)