Add integration tests for training

Since we will be redoing the code, we need an integration test that will be run each time the code changes. This should train a model on a few images and during a few epochs, and compare the model obtained to an expected model.