One possible approach would be to use forward hooks and to compare the intermediate outputs of each layer. This post gives you an example on how to register the hook and how to print it.
You could use different dicts or store the desired outputs using different keys etc.
1 Like