Training scripts need refactoring

[eelcovdw]

• Currently our training scripts are a bit of a mess, because we combine multiple models and different diagnostics (tensorboard etc).
• Gonçalo suggests pytorch lightning (https://github.com/PyTorchLightning/pytorch-lightning). Seems really nice
• Suggestion Wilker for experiment management: wandb. This might not be compatible with lightning.
• Tensorboard could be more flexible: Let the Trainer object decide what to save during training for each model, not the train.py script.
• Why are we using TensorboardX and not pytorch tensorboard?