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

This reverts commit f690d8b1f86de9e2ba8fdda413c9d8264b169980.
---
 torchgeo/datamodules/agrifieldnet.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/torchgeo/datamodules/agrifieldnet.py b/torchgeo/datamodules/agrifieldnet.py
index 787a558ecf4..d1e8db8560b 100644
--- a/torchgeo/datamodules/agrifieldnet.py
+++ b/torchgeo/datamodules/agrifieldnet.py
@@ -7,6 +7,7 @@
 
 import kornia.augmentation as K
 import torch
+from kornia.constants import DataKey, Resample
 
 from ..datasets import AgriFieldNet, random_bbox_assignment
 from ..samplers import GridGeoSampler, RandomBatchGeoSampler
@@ -54,6 +55,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}
+            },
         )
 
     def setup(self, stage: str) -> None: