Improve saving/loading of torch models

[hrzn]

Currently save_model() saves a complete pickle of self (using torch.save()), and in addition it also saves a Pytorch Lightning checkpoint. This is because the torch-save does not save correctly the Trainer object of PTL. But as a result we are saving two files and saving parameters twice. We can probably improve this.