Fix mypy error (#1264)

TorchIO-project · Jan 30, 2025 · ca79007 · ca79007
1 parent 18bc012
commit ca79007
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/torchio/data/inference/aggregator.py b/src/torchio/data/inference/aggregator.py
@@ -104,7 +104,7 @@ def _initialize_avgmask_tensor(self, batch: torch.Tensor) -> None:
         )
 
     @staticmethod
-    def _get_hann_window(patch_size):
+    def _get_hann_window(patch_size) -> torch.Tensor:
         hann_window_3d = torch.as_tensor([1])
         # create a n-dim hann window
         for spatial_dim, size in enumerate(patch_size):
@@ -211,6 +211,7 @@ def add_batch(
                     j_ini:j_fin,
                     k_ini:k_fin,
                 ] += patch
+                assert self._hann_window is not None
                 self._avgmask_tensor[
                     :,
                     i_ini:i_fin,