Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
save_output = SaveOutput()
hook_handles = []
for layer in model.modules():
if isinstance(layer, torch.nn.modules.conv.Conv2d):
handle = layer.register_forward_hook(save_output)
hook_handles.append(handle)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment