Simple Neural Network to Approximate a Quadratic Function

Question

I am currently trying to train a neural network to approximate a quadratic function. The code below does a reasonably good job of doing so:

# Imports
import n