Replace module backward hook with tensor hook

cybertronai / autograd-hacks

The Unlicense

153 stars 32 forks source link

Replace module backward hook with tensor hook #1

Open yaroslavvb opened 5 years ago

yaroslavvb commented 5 years ago

Module backward hooks may get deprecated eventually, can replace them with tensor hook as follows.

def tensor_hook_adder(module, input, output):
    def tensor_backwards(grad):
        setattr(module, 'backprops', grad)
    output.register_hook(tensor_backwards)

layer.register_forward_hook(tensor_hook_adder)