I tried you code with learning_rate=1e-5, seems to be updating the loss and model weights as well.
learning_rate=1e-5