[ONNX] Add export of prim::data (#45747)

Summary: Add export of prim::data Pull Request resolved: #45747 Reviewed By: bdhirsh Differential Revision: D24280334 Pulled By: bzinodev fbshipit-source-id: d21eda84eaba9e690852a72c0e63cbb40eae89bc
pytorch · Nov 5, 2020 · 878032d · 878032d
1 parent 192b296
commit 878032d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -430,6 +430,16 @@ def forward(self, input):
         m1 = torch.randn(3, 4, 5, 6, 7)
         self.run_test(MyModel(), m1)
 
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_data(self):
+        class Data(torch.jit.ScriptModule):
+            @torch.jit.script_method
+            def forward(self, x):
+                return x.new_zeros(x.data.size())
+
+        x = torch.randn(3, 4)
+        self.run_test(Data(), x)
+
     @skipIfUnsupportedMinOpsetVersion(11)
     @disableScriptTest()  # Need type inference
     def test_index_mask_nd(self):

diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -2193,6 +2193,8 @@ def log2(g, self):
 def prim_shape(g, self):
     return g.op('Shape', self)
 
+def prim_data(g, self):
+    return self
 
 @parse_args('v', 'i')
 def one_hot(g, self, num_classes):