def METHOD_NAME(self, x):
    """
    :param x: Input data
    :return: output - mnist digit label for the input image
    """
    batch_size = x.size()[0]
    # (b, 1, 28, 28) -> (b, 1*28*28)
    x = x.view(batch_size, -1)
    # layer 1 (b, 1*28*28) -> (b, 128)
    x = self.layer_1(x)
    x = torch.relu(x)
    # layer 2 (b, 128) -> (b, 256)
    x = self.layer_2(x)
    x = torch.relu(x)
    # layer 3 (b, 256) -> (b, 10)
    x = self.layer_3(x)
    # probability distribution over labels
    x = torch.log_softmax(x, dim=1)
    return x