GradientMap usage to train model, all output converge on my first X value

Question

Growing past using model.fit() to train my model, I wanted to use a GradientMap to implement the training of a basic model that tried to produce the square of a value.

<