Feature: Model + logs saving at the end of training (using safetensors or pytorch)

Make sure all logs are saved at the end of training (see issue #23, as wandb and tensorboard should save stuff automatically anyway).

Save the model using either safetensors or pytorch at the end.