removed TODO in rewrite check

apache · Dec 23, 2020 · bc5e3ec · bc5e3ec
1 parent 8029644
commit bc5e3ec
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/python/frontend/pytorch/test_object_detection.py b/tests/python/frontend/pytorch/test_object_detection.py
@@ -26,8 +26,7 @@
 
 from tvm import relay
 from tvm.runtime.vm import VirtualMachine
-from tvm.relay.dataflow_pattern import rewrite
-from tvm.relay.frontend.pytorch_utils import NMSRewrite
+from tvm.relay.frontend.pytorch_utils import rewrite_nms_to_batched_nms
 from tvm.contrib.download import download
 
 
@@ -142,8 +141,7 @@ def count_valid_scores(scores):
     )
 
     before = mod["main"]
-    after = rewrite(NMSRewrite(), before)
-    # TODO(masahi): Is there a better way to test if the desired rewrite has happened?
+    after = rewrite_nms_to_batched_nms(before)
     assert not tvm.ir.structural_equal(after, before)
 
     mod["main"] = after