Enhancement: Allow both torch.nn.Module as well as HF models

This should most likely be a switch, or perhaps automatic detection.

Allow the training loop for both torch.nn.Module models, as well as models from HF which come with their own set of custom stuff.

Should not affect the training loop itself much, but there will be a couple of changes needed, especially for saving the model, as well as feeding data to its forward() method. Probably much more than that.