diff --git a/docs/op_matrix.md b/docs/op_matrix.md index 0b0a9173..6dec3cd5 100644 --- a/docs/op_matrix.md +++ b/docs/op_matrix.md @@ -94,6 +94,8 @@ Operators that are implemented in Python | `aten::hardtanh_` | | | `aten::im2col` | only 4-D input tensors (batched image-like tensors) are supported | | `aten::index` | Multiple indices for aten::index is not supported | +| `aten::index_put` | aten::index_put_ with accumulate=True is not supported | +| `aten::index_put_` | aten::index_put_ with accumulate=True is not supported | | `aten::index_select` | | | `aten::instance_norm` | | | `aten::layer_norm` | | diff --git a/tinynn/converter/operators/torch/aten.py b/tinynn/converter/operators/torch/aten.py index 685c2d7a..d96fbc63 100644 --- a/tinynn/converter/operators/torch/aten.py +++ b/tinynn/converter/operators/torch/aten.py @@ -3077,7 +3077,10 @@ def parse(self, node, attrs, args, graph_converter): input_names, self.input_tensors[1], graph_converter=graph_converter, non_existent_as_buffer=True ) else: - indices_tensors = [self.find_or_create_input(1, graph_converter)] + if type(self.input_tensors[1]) in (tuple, list): + indices_tensors = [self.create_attr_tensor(x) for x in self.input_tensors[1]] + else: + indices_tensors = [self.find_or_create_input(1, graph_converter)] dim = input_tensor.tensor.ndim