Just as a note, you should not do:
model.w_max.grad.data += w_max1_grad

But
with torch.no_grad():
    model.w_max.grad += w_max1_grad

The first one can hide very subtle bugs that will make the code run but the computed gradients will be wrong.