From 50bc85d266d04e736eb095f46d6b82cfc55c32c2 Mon Sep 17 00:00:00 2001
From: whs <wanghaoshuang@baidu.com>
Date: Tue, 19 May 2020 15:55:01 +0800
Subject: [PATCH] Fix eval_run in pruning script (#724)

---
 slim/prune/prune.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/slim/prune/prune.py b/slim/prune/prune.py
index 6f6dcc4ca..2b7815059 100644
--- a/slim/prune/prune.py
+++ b/slim/prune/prune.py
@@ -307,6 +307,9 @@ def main():
 
             if FLAGS.eval:
                 # evaluation
+                resolution = None
+                if 'Mask' in cfg.architecture:
+                    resolution = model.mask_head.resolution
                 results = eval_run(
                     exe,
                     compiled_eval_prog,
@@ -314,10 +317,8 @@ def main():
                     eval_keys,
                     eval_values,
                     eval_cls,
-                    cfg=cfg)
-                resolution = None
-                if 'mask' in results[0]:
-                    resolution = model.mask_head.resolution
+                    cfg=cfg,
+                    resolution=resolution)
                 box_ap_stats = eval_results(
                     results,
                     cfg.metric,
-- 
GitLab