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

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.

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.

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.

Yep, reviewing this now...