More flexible neural network training

[ThibeauWouters]

Currently, fiesta only has an MLP neural network implemented which uses the ReLU activation function, and the learning rate is fixed. It would be to generalize the training code to make it more flexible so that users can play around with this.

Some ideas:

• Add a learning rate scheduler from optax. This is useful since training now plateaus very quickly and we are losing a lot of potential by not varying the learning rate over time.
• Add other activation functions, make sure they can be saved and loaded without errors
• Add other neural network architectures. Currently, there is only the MLP, of which the layer sizes can be chosen by the user. More advanced architectures could improve the performance.