Add a unit test to check that we can overfit a single example

[JohnGiorgi]

Ideally, the unit test for our model would check that the model can memorize a single training example. This is technically possible by providing the argument metric_to_check to ensure_model_can_train_save_and_load. However, this is reporting an error that the metric cannot be found. I am tabling this for now, but it should be revisited. Opening this issue so I don't forget.