Add dtype field to RasterDataset

torchgeo/datasets/chesapeake.py

@@ -40,6 +40,7 @@ class Chesapeake(RasterDataset, abc.ABC):
     """
 
     is_image = False
+    dtype = torch.long
 
     # subclasses use the 13 class cmap by default
     cmap = {

torchgeo/datasets/geo.py

@@ -296,6 +296,18 @@ class RasterDataset(GeoDataset):
     #: Color map for the dataset, used for plotting
     cmap: dict[int, tuple[int, int, int, int]] = {}
 
+    @property
+    def dtype(self) -> torch.dtype:
+        """The dtype of the dataset (overrides the dtype of the data file via a cast).
+
+        Returns:
+            the dtype of the dataset
+        """
+        if self.is_image:
+            return torch.float32
+        else:
+            return torch.long
+
     def __init__(
         self,
         root: str = "data",
@@ -429,10 +441,12 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
             data = self._merge_files(filepaths, query, self.band_indexes)
 
         sample = {"crs": self.crs, "bbox": query}
+
+        data = data.to(self.dtype)
         if self.is_image:
-            sample["image"] = data.float()
+            sample["image"] = data
         else:
-            sample["mask"] = data.long()
+            sample["mask"] = data
 
         if self.transforms is not None:
             sample = self.transforms(sample)