torch.sign isn’t technically breaking the computation graph, but will return a zero gradient.
I think you could use a sigmoid as a smooth approximation for values in [0, 1].
tanh would approximate the sign method or alternatively you could also use
x / (torch.sqrt(x*x + delta)).