The non-differentiability is from the discontinuity of the function. So you would have to relax (smoothen) the expression.
For non-NaN/inf, where(x > eps, z, y) is (x - eps > 0) * (z - y) + y which you could change to  torch.sigmoid(x - eps) * T) * (z - y) + y for some tuneable float  t.
Best regards
Thomas