fix kthvalue dim init

Xilinx · Dec 21, 2023 · f6bd017 · f6bd017
1 parent 06b810f
commit f6bd017
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/brevitas/utils/torch_utils.py b/src/brevitas/utils/torch_utils.py
@@ -64,7 +64,13 @@ def kthvalue(x, k, dim=None, keepdim=False, out=None) -> torch.Tensor:
     if (dtype == torch.float16 and 'cpu' in device) or \
         (dtype == torch.bfloat16 and 'cuda' in device):
         x = x.type(torch.float32)
-    x, indices = torch.kthvalue(x, k, dim=dim, keepdim=keepdim)
+
+    # PyTorch specify None as default for `dim` but it breaks if we specifically pass None
+    if dim is not None:
+        x, indices = torch.kthvalue(x, k, dim=dim, keepdim=keepdim)
+    else:
+        x, indices = torch.kthvalue(x, k, keepdim=keepdim)
+
     if x.dtype != dtype:
         x = x.type(dtype)
     return (x, indices)