From a14bedcbc8afae7d131eaaa0e467e7b9af06b1f6 Mon Sep 17 00:00:00 2001 From: shukun-ziqiangxu Date: Thu, 6 Jul 2023 00:14:56 +0800 Subject: [PATCH] [Frontend][pytorch]Add copy_ op to pytorch frontend --- python/tvm/relay/frontend/pytorch.py | 5 +++++ tests/python/frontend/pytorch/test_forward.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 5e4d75599613..859f71bab4e6 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -1689,6 +1689,10 @@ def clone(self, inputs, input_types): data = inputs[0] return _op.tensor.copy(data) + def copy_(self, inputs, input_types): + data = inputs[1] + return _op.tensor.copy(data) + def log_softmax(self, inputs, input_types): data = inputs[0] axis = int(inputs[1]) @@ -3841,6 +3845,7 @@ def create_convert_map(self): "aten::reshape": self.reshape, "aten::reshape_as": self.reshape_as, "aten::clone": self.clone, + "aten::copy_": self.copy_, "aten::log_softmax": self.log_softmax, "aten::sigmoid": self.sigmoid, "aten::softplus": self.softplus, diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 83930d1ea80b..29d12d8a237a 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -1707,6 +1707,22 @@ def forward(self, *args): verify_model(Clone1().float().eval(), input_data=input_data) +@tvm.testing.uses_gpu +def test_forward_copy_(): + """test_forward_copy_""" + torch.set_grad_enabled(False) + input_shape = [10] + + class Copy_(Module): + def forward(self, input_data): + r = torch.zeros(input_data.shape).float() + r.copy_(input_data) + return r + + input_data = torch.rand(input_shape).float() + verify_model(Copy_().float().eval(), input_data=input_data) + + @tvm.testing.uses_gpu def test_forward_gather(): """test_forward_gather"""