Finally finished my project. I expected this project to only take like 5 hours max, but it took more than double that! I have learnt a lot about pygame and pytorch in this project. My final additons were a databox that showed ML metrics like epochs, and added a formula to the line on the graph.