Different loss functions for training and validation
I am having a strange problem where I am unable to use different loss functions for pytorch lightning. So, if I use something like this, it is completely fine:
class SegmentationModel(pl.LightningModule):
    def __init__(self, hparams: dict):
        self.lossfn = GeneralizedDiceLoss()

    def training_step(self, batch, batch_idx):
        inputs, targets = self.prepare_batch(batch, output_folders=False)
        logits = self.model(inputs)
        loss_val = self.lossfn(logits, targets)
        return {'loss': loss_val}

    def validation_step(self, batch, batch_idx):
        inputs, targets = self.prepare_batch(batch)
        logits = self.model(inputs)
        loss_val = self.lossfn(logits, targets)
        return {'val_loss': loss_val}


Now, as soon as I want to use a different function (even if it is the same function!), things do not work. In particular, the validation loss does not move. So, if I do something like this in my class definition:
class SegmentationModel(pl.LightningModule):
    def __init__(self, hparams: dict):
        self.train_lossfn = GeneralizedDiceLoss()
        self.val_lossfn = GeneralizedDiceLoss()  # same loss!

    def training_step(self, batch, batch_idx):
        inputs, targets = self.prepare_batch(batch, output_folders=False)
        logits = self.model(inputs)
        loss_val = self.train_lossfn(logits, targets)  # use train loss
        return {'loss': loss_val}

    def validation_step(self, batch, batch_idx):
        inputs, targets = self.prepare_batch(batch)
        logits = self.model(inputs)
        loss_val = self.val_lossfn(logits, targets)  # use validation loss
        return {'val_loss': loss_val}


This makes the model validation completely stuck and it seems the model does not generalise at all. I have a feeling that both the loss functions get called in sequence because the output shows something like:
 | Name         | Type                | Params
-----------------------------------------------------
0 | model        | UNet2D              | 2 M
1 | train_lossfn | GeneralizedDiceLoss | 0
2 | val_lossfn   | GeneralizedDiceLoss | 0


Perhaps I am using this wrong and I probably need to detach the validation loss function but is there a way to achieve this?