This version is not very memory efficient. I trained a  LISTA using the following softshrink :
def softshrink(x, lambd):
    return nn.functional.relu(x - lambd) - nn.functional.relu(-x - lambd)

it saved me 40% memory