Thank you! Do you know how to display the value of the loss function after it's done training? It shows it while the optimization is ongoing, but then the display disappears. Oh, never mind, I figured it out.