Address comments

Thiago Crepaldi · Thiago Crepaldi · commit 0455d360cc12 · 2022-04-28T11:32:11.000-07:00
diff --git a/test/onnx/test_onnx_export.py b/test/onnx/test_onnx_export.py
@@ -7,6 +7,7 @@
 
 import onnx
 import torch
+from typing import Callable
 
 from test_pytorch_common import TestCase
 from torch.onnx.symbolic_helper import _onnx_unsupported
@@ -62,38 +63,23 @@ def break_is_registered_op_api(opname, domain, version):
         self.assertAtenOp(onnx_model, "clamp", "Tensor")
 
 
-    def _helper_test_to_(self, cast_fn):
-        """Run a generic test with the specified cast_fn
+    def _helper_test_to_(self, cast_fn: Callable[[torch.Tensor], torch.Tensor]):
+        """Helper to test aten::to(device) variants
 
-        `cast_fn` will be converted into a `torch.jit.script` function
-        and it is meant to wrap a `aten::to` during export to prevent
-        hard-coded devices
+        `cast_fn` is converted into a `torch.jit.script`. It wraps `aten::to`
+        during export to preventing the devices to be hard-coded.
 
-        `cast_fn` signature is:
-            def cast_fn(input: torch.Tensor) -> torch.Tensor
+        Needed by detectron2 after https://github.com/facebookresearch/detectron2/pull/4132/
         """
-        cast_fn = torch.jit.script_if_tracing(cast_fn)
-
-        class MyModel(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.conv1 = torch.nn.Conv2d(3, 20, 5)
-                self.conv2 = torch.nn.Conv2d(20, 20, 5)
-
-            def forward(self, x):
-                x = cast_fn(x)
-                x = torch.nn.functional.relu(self.conv1(x))
-                return torch.nn.functional.relu(self.conv2(x))
+        cast_fn = torch.jit.script(cast_fn)
 
         f = io.BytesIO()
-        model = MyModel()
         x = torch.zeros([1, 3, 32, 32])
-        torch.onnx.export(model, (x,), f,
-                          operator_export_type=OperatorExportTypes.ONNX)
+        torch.onnx.export(cast_fn, (x,), f)
         onnx_model = onnx.load_from_string(f.getvalue())
         for n in onnx_model.graph.node:
-            assert n.op_type != "To"
-            assert n.op_type != "Cast"
+            self.assertNotEqual(n.op_type, "To")
+            self.assertNotEqual(n.op_type, "Cast")
 
     def test_to__cpu_string(self):
         def cast_cpu_string(src: torch.Tensor) -> torch.Tensor:
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -2249,33 +2249,29 @@ def topk(g, self, k, dim, largest, sorted, out=None):
 
 def to(g, self, *args):
 
-    def is_aten_to_device_only(*args):
+    def is_aten_to_device_only(args):
         if len(args) == 4:
             # aten::to(Tensor, Device, bool, bool, memory_format)
-            if args[0].node().kind() == "prim::device" or \
-                    args[0].type().isSubtypeOf(ListType.ofInts()):
-                return True
-            if sym_help._is_value(args[0]) and \
-                    args[0].node().kind() == "onnx::Constant" and \
-                    isinstance(args[0].node()["value"], str):
-                return True
+            return args[0].node().kind() == "prim::device" or \
+                args[0].type().isSubtypeOf(ListType.ofInts()) or \
+                (sym_help._is_value(args[0]) and
+                    args[0].node().kind() == "onnx::Constant" and
+                    isinstance(args[0].node()["value"], str))
         elif len(args) == 5:
             # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format)
+            # When dtype is None, this is a aten::to(device) call
             dtype = sym_help._get_const(args[1], "i", "dtype")
-            if not dtype:
-                # When dtype is None, this is a aten::to(device) call
-                return True
-        elif len(args) >= 6 and len(args) <= 7:
+            return dtype is None
+        elif len(args) in (6, 7):
             # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format) -> Tensor
             # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format) -> Tensor
+            # When dtype is None, this is a aten::to(device) call
             dtype = sym_help._get_const(args[0], "i", "dtype")
-            if not dtype:
-                # When dtype is None, this is a aten::to(device) call
-                return True
+            return dtype is None
         return False
 
     # ONNX doesn't have a concept of a device, so we ignore device-only casts
-    if is_aten_to_device_only(*args):
+    if is_aten_to_device_only(args):
         return self
 
     if len(args) == 4:
