Solved by InnovArul in post #2 
              
              
                I am not sure about your exact use case. In my understanding, if you are not using any forward or backward hooks, you can always define additional methods inside nn.Module. 
Note that the hooks would only work if model(input) is called, which in turn processes the hooks internally before and after c…