try to fix mypy

torchgeo/datamodules/digital_typhoon.py

@@ -5,8 +5,6 @@
 
 from typing import Any, Union
 
-from torch.utils.data import Dataset
-
 from ..datasets import DigitalTyphoonAnalysis
 from .geo import NonGeoDataModule
 from .utils import group_shuffle_split
@@ -43,8 +41,8 @@ def __init__(
         self.split_by = split_by
 
     def split_dataset(
-        self, dataset: Dataset
-    ) -> tuple[list[dict[str, Union[str, list]]]]:
+        self, dataset: DigitalTyphoonAnalysis
+    ) -> tuple[list[int, dict[str, Union[str, list[int]]]]]:
         """Split dataset into two parts.
 
         Args:

torchgeo/datasets/digital_typhoon.py

@@ -279,6 +279,15 @@ def _load_image(self, sample_df: pd.DataFrame) -> Tensor:
         """
 
         def load_image_tensor(id: str, filepath: str) -> Tensor:
+            """Load a single image tensor from a h5 file.
+
+            Args:
+                id: typhoon id
+                filepath: path to the h5 file
+
+            Returns:
+                image tensor
+            """
             full_path = os.path.join(self.root, self.data_root, "image", id, filepath)
             with h5py.File(full_path, "r") as h5f:
                 # tensor with added channel dimension