Cast uint16 tensors to int32, leave float32 tensors as they are (micr…

…osoft#502)
yichiac · Apr 11, 2022 · ba820e3 · ba820e3
1 parent 110c919
commit ba820e3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchgeo/datasets/sen12ms.py b/torchgeo/datasets/sen12ms.py
@@ -266,7 +266,9 @@ def _load_raster(self, filename: str, source: str) -> Tensor:
                 "{}_{}_{}_{}_{}".format(*parts),
             )
         ) as f:
-            array = f.read().astype(np.int32)
+            array = f.read()
+            if array.dtype == np.uint16:
+                array = array.astype(np.int32)
             tensor = torch.from_numpy(array)
             return tensor