From 9acf85940e7c02636e383788be208f8e17e5351f Mon Sep 17 00:00:00 2001
From: "Adam J. Stewart" <ajstewart426@gmail.com>
Date: Fri, 15 Mar 2024 13:40:13 +0100
Subject: [PATCH] Revert "extra_args not needed"

This reverts commit 859f24ef151607a23d82afc985022d5630bd95ca.
---
 torchgeo/datamodules/sentinel2cdl.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/torchgeo/datamodules/sentinel2cdl.py b/torchgeo/datamodules/sentinel2cdl.py
index 5acdac832fa..ac79ebabfc8 100644
--- 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(