def METHOD_NAME(self, stage=None):
    """
    Downloads the data, parse it and split the data into train, test, validation data
    :param stage: Stage - training or testing
    """
    self.df_train = datasets.MNIST(
        "dataset", download=True, train=True, transform=self.transform
    )
    self.df_train, self.df_val = random_split(self.df_train, [55000, 5000])
    self.df_test = datasets.MNIST(
        "dataset", download=True, train=False, transform=self.transform
    )