Revert "extra_args not needed"

This reverts commit 859f24e.
microsoft · Mar 15, 2024 · 9acf859 · 9acf859
1 parent 4c05652
commit 9acf859
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torchgeo/datamodules/sentinel2cdl.py b/torchgeo/datamodules/sentinel2cdl.py
@@ -7,6 +7,7 @@
 
 import kornia.augmentation as K
 import torch
+from kornia.constants import DataKey, Resample
 
 from ..datasets import CDL, Sentinel2, random_grid_cell_assignment
 from ..samplers import GridGeoSampler, RandomBatchGeoSampler
@@ -67,6 +68,9 @@ def __init__(
             K.RandomVerticalFlip(p=0.5),
             K.RandomHorizontalFlip(p=0.5),
             data_keys=["image", "mask"],
+            extra_args={
+                DataKey.MASK: {"resample": Resample.NEAREST, "align_corners": None}
+            },
         )
 
         self.aug = AugmentationSequential(