How can catch the animation for much deep network like ResNet using Deepreplay?

Hey How could assign this method on ResNet? I tried but it throws out the following error:
model.summary()
==================================================================================================
_________________________________________________________________________________________________
conv2d_36 (Conv2D)              (None, 1, 1, 512)    2359808     activation_32[0][0]              
__________________________________________________________________________________________________
add_16 (Add)                    (None, 1, 1, 512)    0           add_15[0][0]                     
                                                                 conv2d_36[0][0]                  
__________________________________________________________________________________________________
batch_normalization_33 (BatchNo (None, 1, 1, 512)    2048        add_16[0][0]                     
__________________________________________________________________________________________________
activation_33 (Activation)      (None, 1, 1, 512)    0           batch_normalization_33[0][0]     
__________________________________________________________________________________________________
average_pooling2d_1 (AveragePoo (None, 1, 1, 512)    0           activation_33[0][0]              
__________________________________________________________________________________________________
flatten_1 (Flatten)             (None, 512)          0           average_pooling2d_1[0][0]        
__________________________________________________________________________________________________
dense_1 (Dense)                 (None, 10)           5130        flatten_1[0][0]                  
==================================================================================================
from keras.callbacks import TensorBoard
tensorboard = TensorBoard(log_dir='./logs_ResNet34', histogram_freq=1,
                          write_graph=True, write_images=True)
# define model
model.fit(X_train, Y_train,
          batch_size=batch_size,
          epochs=nb_epoch,
          validation_data=(X_test, Y_test),
          shuffle=True,
          callbacks=[lr_reducer, early_stopper, csv_logger, tensorboard, replaydata])

filename = 'hyperparms_in_action.h5'
group_name = 'part2'

replaydata = ReplayData(X_train, Y_train, filename=filename, group_name=group_name, model=model)

fig, ax = plt.subplots(1, 1, figsize=(5, 5))
fs = replay.build_feature_space(ax, layer_name='dense_1') #AssertionError: Only layers with 2-dimensional outputs are supported!
#fs = replay.build_decision_boundary(ax_fs)               #NameError: name 'ax_fs' is not defined
fs.plot(epoch=60).savefig('feature_space_epoch60.png', dpi=120)
fs.animate().save('feature_space_animation.mp4', dpi=120, fps=5)
How can I fix that?
dvgodoy / deepreplay

How can catch the animation for much deep network like ResNet using Deepreplay? #18
