Dumb fix

cifar-10-python.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53c8b1a8d9b386502f69a87ab4d636d25a893ba9c434be754980ec283a5bd227
+size 3768320

main.py

@@ -28,7 +28,7 @@ def main(model_config=None):
         "nrow": 8
     }
 
-    if modelConfig['state'] == train:
+    if modelConfig['state'] == 'train':
         train(modelConfig)
     else:
         eval(modelConfig)

train.py

@@ -15,7 +15,7 @@ def train(modelConfig: Dict):
     device = torch.device(modelConfig['device'])
 
     dataset = CIFAR10(
-        "./",train=True,transform=transforms.Compose([
+        "./",train=True,download=True,transform=transforms.Compose([
             transforms.RandomHorizontalFlip(),
             transforms.ToTensor(),
             transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
@@ -60,7 +60,7 @@ def train(modelConfig: Dict):
                     "LR": optimizer.state_dict()['param_groups'][0]["lr"]
                 })
         warmupScheduler.step()
-        torch.save(net_model,os.path.join(modelConfig['checkpoint_dir'] + f"ckpt_{epoch}_.pth"))
+        torch.save(net_model,os.path.join(modelConfig['checkpoint_dir'] + f"ckpt_{epoch}.pth"))
 
 
 def eval(modelConfig:Dict):