Challenge: Optimizers

Explore the application of SGD and AdaBelief optimizers on a convolutional neural network using JAX and Flax. Learn to initialize these optimizers, monitor loss variations, and evaluate which optimizer yields better training performance in this scenario.