Solved by antigol in post #3 
              
              
                It worked! Thanks 
class SmoothStep(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        y = torch.zeros_like(x)
        m = (x > 0.0)
        y[m] = (-1/x[m]).exp()
        return y

    @staticmethod
    def backward(ctx, dy):
        x, = c…