Add sum(idx is not None for idx in indices) == 1 check to fix incorre…

…ct mapping of multi-indexing in TVM logic (#51)
tenstorrent · Dec 5, 2024 · 38d2522 · 38d2522
1 parent 24e1c49
commit 38d2522
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -2832,7 +2832,7 @@ def index(self, inputs, input_types):
 
             return res
 
-        elif len(_infer_shape(data)) > 2 :
+        elif len(_infer_shape(data)) > 2 and sum(idx is not None for idx in indices) == 1:
             axis = None
             index_expr = None
             for i, idx in enumerate(indices):