    def validation_step(self, val_batch, batch_idx):
        """
        Performs validation of data in batches

        :param val_batch: Batch data
        :param batch_idx: Batch indices

        :return: output - valid step loss
        """
        x, y = val_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        return {"val_step_loss": loss}
Performs validation of data in batches

:param val_batch: Batch data
:param batch_idx: Batch indices

:return: output - valid step loss
tests/integration/frameworks/mlflow/MNIST/mnist_autolog_example.py