def prepare_data(self):
# download only
MNIST(self.data_dir,
train=True,
download=True,
normalize=(0.1307, 0.3081))
MNIST(self.data_dir,
train=False,
download=True,
normalize=(0.1307, 0.3081))
def setup(self, stage: Optional[str] = None):
# Assign train/val datasets for use in dataloaders
# TODO: need to split using random_split once updated to torch >= 1.6
if stage == "fit" or stage is None:
self.mnist_train = MNIST(self.data_dir,
train=True,
normalize=(0.1307, 0.3081))
# Assign test dataset for use in dataloader(s)
if stage == "test" or stage is None:
self.mnist_test = MNIST(self.data_dir,
train=False,
normalize=(0.1307, 0.3081))
def train_dataloader(self):
return DataLoader(MNIST(
train=True,
download=True,
),
batch_size=128,
num_workers=1)