    def test_step(self, test_batch, batch_idx):
        """
        Performs test and computes the accuracy of the model

        :param test_batch: Batch data
        :param batch_idx: Batch indices

        :return: output - Testing accuracy
        """
        x, y = test_batch
        output = self.forward(x)
        _, y_hat = torch.max(output, dim=1)
        test_acc = accuracy(y_hat.cpu(), y.cpu())
        return {"test_acc": test_acc}
Performs test and computes the accuracy of the model

:param test_batch: Batch data
:param batch_idx: Batch indices

:return: output - Testing accuracy
tests/integration/frameworks/mlflow/MNIST/mnist_autolog_example.py