Move loss to lightning module/separate module

[djdameln]

Switch to a new convention for where the loss function is defined.

From the discussion:

The torch_model might not be the appropriate place to define the loss function. The torch model class does not contain any training logic, so it never actually uses the loss. In those models that have the loss function as an attribute of the torch model, self.loss is never accessed from within the torch model (e.g. STFPM).

My proposal would be to make the loss an attribute of the lightning module instead. In that case we can call self.loss(...) instead of self.model.loss(...) in the training_step method. If the total loss consists of a simple combination of multiple loss functions (e.g. DRAEM), we can perform the computation within the training_step, while if more advanced logic is involved (e.g. GANomaly), we can define a helper class in loss.py.

This way we would keep the responsibilities of the different modules separated:

• torch_model.py: Architecture of the model
• lightning_model.py: Training and validation logic
• loss.py: Helper for advanced loss computation