fix test

apache · Dec 24, 2020 · 35a177b · 35a177b
1 parent bc5e3ec
commit 35a177b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/python/frontend/pytorch/test_object_detection.py b/tests/python/frontend/pytorch/test_object_detection.py
@@ -126,7 +126,7 @@ def compile_and_run_vm(mod, params, data_np):
 
     def count_valid_scores(scores):
         num_valid_scores = 0
-        for score in pt_scores:
+        for score in scores:
             if score >= score_threshold:
                 num_valid_scores += 1
             else:
@@ -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(