Add tests for overfitting

[l3th4l]

We were able to overfit on a batch_size of 25 and num_batches=1 with the train_with_env.py, up till #3
We need a pytest script which makes sure that any future training ensures that we can achieve similar performance with --use_error_mask False