typing

Xilinx · Dec 21, 2023 · edba342 · edba342
1 parent e996018
commit edba342
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/utils/torch_utils.py b/src/brevitas/utils/torch_utils.py
@@ -54,7 +54,7 @@ def kthvalue(
         k: int,
         dim: Optional[int] = None,
         keepdim: bool = False,
-        out: Optional[Tuple] = None) -> torch.Tensor:
+        out: Optional[Tuple[torch.Tensor, torch.LongTensor]] = None) -> torch.Tensor:
     # As of torch 2.1, there is no kthvalue implementation:
     # - In CPU for float16
     # - In GPU for bfloat16