datamodules: Switch to kornia AugmentationSequential

[ashnair1]

Switch to kornia's AugmentationSequential in datamodules submodule

Closes #1432

TODO (once kornia 0.7.4 is released):

• Remove manual setting of keepdim attribute
• Bump min version of kornia

[ashnair1]

TODO: Tutorials need to be updated once we fully switch to kornia's AugmentationSequential

[adamjstewart]

Still confused by some of the dataset changes, but all datamodule changes now look good. I'll see if I can figure out why the tests are failing.

[adamjstewart]

But tests are failing for most GeoDataModules.

I wish it were this simple... Note that some failing tests are for NonGeoDataModules (ssl4eo_l_benchmark_*).

Also note that some tests fail depending on batch size. For example, gid15 passes with batch_size=1, but fails if batch_size=2. But there are other datasets for which both 1 and 2+ pass. For the case of gid15, it seems like this is because the image is 1x1, so this may be a red herring.

Most of the failing tests seem to be because the ground truth mask has a channel dimension, which cross-entropy loss does not like. Let me see if I can fix a few.

[adamjstewart]

The ChesapeakeCVPR dataset is royally screwed up and I'm not touching it, so I'll let you or @calebrob6 solve that one. The problem is that it doesn't use RasterDataset and the mask is sometimes single channel, sometimes multichannel, so you have to be very careful with all dimensions. I'm not even sure our trainer supports multichannel masks.

[adamjstewart]

Is this waiting on Kornia 0.7.4? Trying to decide if this should be in TorchGeo 0.6.0 next week or if we should bump it to a later release.

[robmarkcole]

Kornia 0.7.4. is out :-)

[adamjstewart]

Yep, reviewing this now...

[adamjstewart]

This looks really clean now, thanks for all the hard work on this!

I think the only other place we use our wrapper is in the transforms tutorial? Maybe we can remove that in a separate PR so we don't hold this up.