fix test

apache · Dec 24, 2020 · 07f31f4 · 07f31f4
1 parent bc5e3ec
commit 07f31f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/python/frontend/pytorch/test_object_detection.py b/tests/python/frontend/pytorch/test_object_detection.py
@@ -141,10 +141,10 @@ def count_valid_scores(scores):
     )
 
     before = mod["main"]
-    after = rewrite_nms_to_batched_nms(before)
+    mod = rewrite_nms_to_batched_nms(mod)
+    after = mod["main"]
     assert not tvm.ir.structural_equal(after, before)
 
-    mod["main"] = after
     tvm_res_after_rewrite = compile_and_run_vm(mod, params, data_np)
 
     num_tvm_after_rewrite_valid_scores = count_valid_scores(