Adding aten::unsqueeze_ to PT Frontend (apache#7231)

* Added Ops * Regular * Remove copy * Remove copy * Tests * Black Co-authored-by: Ubuntu <ubuntu@ip-172-31-27-149.us-east-2.compute.internal> Co-authored-by: Ubuntu <ubuntu@ip-172-31-19-34.us-east-2.compute.internal>
masahi · Jan 18, 2021 · bb37f88 · bb37f88
1 parent 0529446
commit bb37f88
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -2117,6 +2117,7 @@ def create_convert_map(self):
             "aten::to": self.to,
             "aten::squeeze": self.squeeze,
             "aten::unsqueeze": self.unsqueeze,
+            "aten::unsqueeze_": self.unsqueeze,
             "aten::cat": self.concatenate,
             "aten::slice": self.slice,
             "aten::split": self.split,

diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
@@ -447,8 +447,16 @@ class Unsqueeze1(Module):
         def forward(self, *args):
             return args[0].unsqueeze(2)
 
+    class Unsqueeze2(Module):
+        def forward(self, *args):
+            _ = args[0].unsqueeze_(2)
+            # Check whether operations after inplace unsqueeze works as expected
+            y = args[0].squeeze(2)
+            return torch.add(y, y)
+
     input_data = torch.rand(input_shape).float()
     verify_model(Unsqueeze1().float().eval(), input_data=input_data)
+    verify_model(Unsqueeze2().float().eval(), input_data=input_data)
 
 
 @tvm.testing.uses_gpu