More typing

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