fix type-hint

strongio · Apr 19, 2024 · b676e2c · b676e2c
1 parent 634a681
commit b676e2c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchcast/internals/utils.py b/torchcast/internals/utils.py
@@ -136,7 +136,7 @@ def ragged_cat(tensors: Sequence[torch.Tensor],
 
 
 @torch.no_grad()
-def true1d_idx(arr: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
+def true1d_idx(arr: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
     if not isinstance(arr, torch.Tensor):
         arr = torch.as_tensor(arr)
     arr = arr.bool()