Solved by mrTsjolder in post #7 
              
              
                This seems to be working, thanks! The only issue is when the inputs do not require gradients (i.e. in the input layer). However, this can easily be resolved by using something like: 
def forward(self, x):
    if x.requires_grad:
         x = x.clone()
         x.register_hook(lambda grad: grad * sel…