Update resnet.py

resnet.py

@@ -184,12 +184,12 @@ class custom_ResNet(pl.LightningModule):
 
         # Assign train/val datasets for use in dataloaders
         if stage == "fit" or stage is None:
-            cifar_full = CIFAR10(self.data_dir, train=True, transform=self.train_transform)
+            cifar_full = CIFAR10(self.data_dir, train=True, download=True, transform=self.train_transform)
             self.cifar_train, self.cifar_val = random_split(cifar_full, [45000, 5000])
 
         # Assign test dataset for use in dataloader(s)
         if stage == "test" or stage is None:
-            self.cifar_test = CIFAR10(self.data_dir, train=False, transform=self.test_transform)
+            self.cifar_test = CIFAR10(self.data_dir, train=False, download=True, transform=self.test_transform)
 
     def train_dataloader(self):
         return DataLoader(self.cifar_train, batch_size=BATCH_SIZE, num_workers=os.cpu_count())