You could store the activations in a dict using forward hooks. I’ve created a small example in this thread.
dict