Added graceful exit if data download errored (#3096)

pytorch · Oct 16, 2023 · 624eb63 · 624eb63
1 parent 253b5f1
commit 624eb63
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/examples/super_resolution/main.py b/examples/super_resolution/main.py
@@ -73,8 +73,12 @@ def __len__(self):
         return len(self.dataset)
 
 
-trainset = torchvision.datasets.Caltech101(root="./data", download=True)
-testset = torchvision.datasets.Caltech101(root="./data", download=False)
+try:
+    trainset = torchvision.datasets.Caltech101(root="./data", download=True)
+    testset = torchvision.datasets.Caltech101(root="./data", download=False)
+except RuntimeError:
+    print("Dataset download problem, exiting without error code")
+    exit(0)
 
 trainset_sr = SRDataset(trainset, scale_factor=opt.upscale_factor, crop_size=opt.crop_size)
 testset_sr = SRDataset(testset, scale_factor=opt.upscale_factor, crop_size=opt.crop_size)