GraphConv now adative for out-of-bound k.

Signed-off-by: Shivam Pandey <pandeyshivam2017robotics@gmail.com>
ShivamPR21 · Jun 30, 2023 · 590e9c6 · 590e9c6
1 parent a1c4832
commit 590e9c6
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/moduleZoo/graphs/grah_conv.py b/moduleZoo/graphs/grah_conv.py
@@ -12,6 +12,8 @@
 def get_graph_features(x: torch.Tensor, idx: torch.Tensor | None = None, k: int | None = None, mode: str = 'local+global') -> torch.Tensor:
     B, n, d = x.size() # [B, n, d]
 
+    k = min(n, k) if idx is None else k
+
     features = knn_features(x, idx, k) # [B, n, k, d]
     x = x.view(B, n, 1, d).repeat(1, 1, k, 1) # [B, n, k, d]