Skip to content

Instantly share code, notes, and snippets.

@sadimanna

sadimanna/fwd_pre_hook.py

Created January 9, 2021 11:14
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save sadimanna/3f9ee1da0f435c1bb28dbe2e2e018f2a to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save sadimanna/3f9ee1da0f435c1bb28dbe2e2e018f2a to your computer and use it in GitHub Desktop.
Download ZIP
Raw
fwd_pre_hook.py
def forward_pre_hook(self,layer_name):
def pre_hook(module, input):
return torch.zeros(input[0].shape,dtype=torch.float,device='cuda:0',requires_grad = True)
return pre_hook
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment