Update train.py

soniamartinot · Jun 11, 2021 · 58bd0f1 · 58bd0f1
1 parent 3118d2c
commit 58bd0f1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -22,7 +22,7 @@
 # Parser parameters
 args = parse_args()
 # Get the cases paths
-cases = list_cases(simu_path, exclude=[])
+cases = list_cases(args.simu_path, exclude=[])
 # Instantiate the selected model and indicate whether it's UNet based or not
 model, unet = instantiate_model(args.model_name)
 # Create train and validation dataloaders
@@ -241,4 +241,4 @@
 
                 # Reset
                 loss_train, ssim_train, mse_train, l1_train = 0, 0, 0, 0
-                a = time.time()
+                a = time.time()