Solved by albanD in post #4 
              
              
                Defining my_hook on the class will make it a class method. And so you need to specify it as self.my_hook. 
Otherwise you can put it inline in the function: 
import torch
import torch.nn as nn
class Net(torch.nn.Module):

    def __init__(self):       
        super(Net, self).__init__()
        self…