fix eval mode without srn (#1889)

* fix base model * fix start time
2021-01-31 22:37:30 +08:00 · 2021-01-31 22:37:30 +08:00 · 2a0c3d4dac
parent e40100c5b7
commit 2a0c3d4dac
1 changed files with 5 additions and 2 deletions
--- a/tools/program.py
+++ b/tools/program.py
@ -326,9 +326,12 @@ def eval(model, valid_dataloader, post_process_class, eval_class):
            if idx >= len(valid_dataloader):
                break
            images = batch[0]
-            others = batch[-4:]
            start = time.time()
+            if "SRN" in str(model.head):
+                others = batch[-4:]
                preds = model(images, others)
+            else:
+                preds = model(images)

            batch = [item.numpy() for item in batch]
            # Obtain usable results from post-processing methods