Add a test for training MNIST on CUDA and CPU that verifies same result as Jax

[sogartar]

The two tests fail for different reasons. Things work OK except the procedure to initialize the optimizer.

[ThomasRaoux]

is there a way to get a test independent of JAX? If you don't have a simple way I'll try to make one later