fix dataloader loading

cyli2019 · Dec 11, 2018 · 1afd340 · 1afd340
1 parent d9439a2
commit 1afd340
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/cifar10/dataloader.py b/cifar10/dataloader.py
@@ -6,10 +6,10 @@
 
 def get_data_loaders(args):
     if args.trainloader and args.testloader:
-        assert os.path.exists(trainloader_path), 'trainloader does not exist'
-        assert os.path.exists(testloader_path), 'testloader does not exist'
-        trainloader = torch.load(trainloader_path)
-        testloader = torch.load(testloader_path)
+        assert os.path.exists(args.trainloader), 'trainloader does not exist'
+        assert os.path.exists(args.testloader), 'testloader does not exist'
+        trainloader = torch.load(args.trainloader)
+        testloader = torch.load(args.testloader)
         return trainloader, testloader
 
     normalize = transforms.Normalize(mean=[x/255.0 for x in [125.3, 123.0, 113.9]],