fix: handle cases in which losses are nan during the training

[gcroci2]

During the training (in the trainer.Trainer.train() method), if the loss is nan (such as when using very few datapoints, like in #528), checkpoint_model is never created, thus the crashes and gives the following error:

UnboundLocalError: local variable 'checkpoint_model' referenced before assignment

This PR adds checks for handling such cases and integration tests.

• I used tmpdir_factory fixture for generating the hdf5 files needed for the checks only once per tests' session.
• I used pytest.mark.parametrize for passing in multiple inputs to the test function.