Add sanity checks

I thing two things would be great:

1. Simple sanity checks (unit tests) that run after every push in the pipeline e.g. a dummy dataset with just on batch of examples. That should also run in the ci. Just to check al stages are working correctly
2. Some units tests we run every week or so to check that the model is still performing well on a know yet small dataset (e.g. movies)