Skip to content

Instantly share code, notes, and snippets.

@cosmic-cortex

cosmic-cortex/registering_hook.py Secret

Created April 29, 2020 06:49
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save cosmic-cortex/3607750134d806a6755064eab45d163f to your computer and use it in GitHub Desktop.
Save cosmic-cortex/3607750134d806a6755064eab45d163f to your computer and use it in GitHub Desktop.
Download ZIP
Raw
registering_hook.py
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