diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index cc9217c9f5f8..37e8e352c6a7 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -361,6 +361,7 @@ def create_convert_map( "mean.dim": self._mean, "prod.default": self._prod, "std.correction": self._std, + "sum.default": self._sum, "sum.dim_IntList": self._sum, "var.correction": self._var, # search @@ -389,6 +390,7 @@ def create_convert_map( "split_with_sizes.default": self._split, "squeeze.default": self._squeeze, "squeeze.dim": self._squeeze, + "stack.default": self._stack, "take.default": self._take, "tile.default": self._tile, "topk.default": self._topk, diff --git a/tests/python/relax/test_from_exported_to_cuda.py b/tests/python/relax/test_from_exported_to_cuda.py index 56ee527caf09..8405f48576d8 100644 --- a/tests/python/relax/test_from_exported_to_cuda.py +++ b/tests/python/relax/test_from_exported_to_cuda.py @@ -512,5 +512,33 @@ def forward(self, x): assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) +@tvm.testing.parametrize_targets("cuda") +def test_stack(target, dev): + class StackModel(nn.Module): + def forward(self, x): + val1 = x[1, 4] + val2 = x[3, 2] + val3 = x[5, 6] + z = torch.stack([val1, val2, val3]) + return z + + torch_module = StackModel().eval() + raw_data = np.random.rand(10, 10, 10).astype("float32") + assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) + + +@tvm.testing.parametrize_targets("cuda") +def test_sum(target, dev): + class SumModel(nn.Module): + def forward(self, x): + new_vec = x[1, 4] + return new_vec.sum() + + torch_module = SumModel().eval() + + raw_data = np.random.rand(10, 10, 10).astype("float32") + assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) + + if __name__ == "__main__": tvm.testing.main()