Solved by nthn_clmnt in post #5 
              
              
                This should work for what you want, all wrapped up as the activation function ‘non_sat_relu’. The backward leak is set to 0.1, but it could of course be whatever you want. 
import torch

class NSReLU(torch.autograd.Function):
    @staticmethod
    def forward(self,x):
        self.neg = x < 0
      …