Fix plotting in datamodules when dataset is a Subset (#2003)

* Fix plotting in datamodules when dataset is a Subset

* Guessing at how to fix isort

torchgeo/datamodules/geo.py

@@ -11,7 +11,7 @@
 from lightning.pytorch import LightningDataModule
 from matplotlib.figure import Figure
 from torch import Tensor
-from torch.utils.data import DataLoader, Dataset, default_collate
+from torch.utils.data import DataLoader, Dataset, Subset, default_collate
 
 from ..datasets import GeoDataset, NonGeoDataset, stack_samples
 from ..samplers import (
@@ -157,6 +157,8 @@ def plot(self, *args: Any, **kwargs: Any) -> Figure | None:
         """
         fig: Figure | None = None
         dataset = self.dataset or self.val_dataset
+        if isinstance(dataset, Subset):
+            dataset = dataset.dataset
         if dataset is not None:
             if hasattr(dataset, "plot"):
                 fig = dataset.plot(*args, **kwargs)