Bug fix: ONNX export full operator (#21669)

BowenBao · facebook-github-bot · commit 3e8dc565bd69 · 2019-06-13T13:20:21.000-07:00
Summary: Fix an obvious bug. Pull Request resolved: #21669 Reviewed By: zrphercule Differential Revision: D15806614 Pulled By: houseroad fbshipit-source-id: d0f6e934252e0057f3dbcc7f160236ee6f4497ac
diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py
@@ -1250,6 +1250,16 @@ def forward(self, x):
         self.run_model_test(TensorFactory(), train=False, input=(x,), batch_size=BATCH_SIZE,
                             use_gpu=False, example_outputs=(torch.ones(x.size()),))
 
+    def test_full_script(self):
+        class FullClass(torch.jit.ScriptModule):
+            @torch.jit.script_method
+            def forward(self, x):
+                return torch.full((4, 5), x, dtype=torch.long)
+
+        x = torch.tensor(12)
+        self.run_model_test(FullClass(), train=False, input=(x,), batch_size=BATCH_SIZE,
+                            use_gpu=False, example_outputs=FullClass()(x))
+
     def test_where_functional(self):
         class WhereFunctional(torch.nn.Module):
             def forward(self, x):
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -1132,8 +1132,8 @@ def ones_like(g, input, dtype, layout, device, pin_memory=False):
 def full(g, sizes, value, dtype, layout, device, pin_memory=False):
     const_value = sym_help._maybe_get_const(value, 't')
     if sym_help._is_value(const_value):
-        tmp = zeros(sizes, dtype, layout, device)
-        return add(tmp, value, g.op("Constant", value_t=torch.tensor(1)))
+        tmp = zeros(g, sizes, dtype, layout, device)
+        return add(g, tmp, value, g.op("Constant", value_t=torch.tensor(1)))
     else:
         dtype = sym_help._get_const(dtype, 'i', 'dtype')
         return g.op("ConstantOfShape", sizes,
