Merge pull request #12 from rsuwa/fix_attribute_error

Fixed to confirm the existence of the attribute "masks"
Shreeyak · Nov 30, 2020 · 78388b4 · 78388b4
2 parents 8d32dc3 + 69d27b0
commit 78388b4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_networks/surface_normals/train.py b/pytorch_networks/surface_normals/train.py
@@ -225,7 +225,7 @@
 if config.train.datasetsTestReal is not None:
     for dataset in config.train.datasetsTestReal:
         if dataset.images:
-            mask_dir = dataset.masks if dataset.masks else ''
+            mask_dir = dataset.masks if hasattr(dataset, 'masks') and dataset.masks else ''
             db = dataloader.SurfaceNormalsDataset(input_dir=dataset.images,
                                                   label_dir=dataset.labels,
                                                   mask_dir=mask_dir,