share steps among test and validation steps

[ItamarKanter]

I'm implementing the same functionality for validation_step and test_step.
Currently, I have implemented it by calling to a shared function (val_and_test_step)

    def val_and_test_step(self, data_batch, batch_nb):
        output = shared_functionality(data_batch, batch_nb)
        return output

    def validation_step(self, data_batch, batch_nb):
        return self.val_and_test_step(data_batch, batch_nb)

    def test_step(self, data_batch, batch_nb):
        return self.val_and_test_step(data_batch, batch_nb) 

Is there a more pythonic way to implement the above?

[tchaton]

Dear @ItamarKanter,

I believe this great and pretty pythonic !

You could do this to make it slightly cleaner.

class Model(LightningModule)

    def common_step(self, batch, batch_idx, stage):
        logits = self(batch[0])
        loss = self.compute_loss(logits, batch[1])
        self.log(f"{state}_loss", loss)
        return loss

    def training_step(self, batch, batch_idx):
        return self.common_step(batch, batch_idx, "train")

    def validation_step(self, batch, batch_idx):
        self.common_step(batch, batch_idx, "val")

    def test_step(self, batch, batch_idx):
        self.common_step(batch, batch_idx, "test")

Best,
T.C