JIT annotation

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