Skip to content

Commit caa76d8

Browse files
xuzhao9facebook-github-bot
authored andcommitted
Fix yolov3 train accuracy test (#2321)
Summary: Setup the correct example inputs for the yolov3 train accuracy test Fixes #2248 Pull Request resolved: #2321 Reviewed By: aaronenyeshi Differential Revision: D58823036 Pulled By: xuzhao9 fbshipit-source-id: d63c069bdf0da6ba496f274510e64d76e3c10f76
1 parent 5254910 commit caa76d8

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

torchbenchmark/models/yolov3/__init__.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,9 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
5454
self.num_epochs = 1
5555
self.train_num_batch = 1
5656
self.prefetch = True
57-
if test == "train":
57+
if test == "eval" or self.dargs.accuracy:
58+
self.model, self.example_inputs = self.prep_eval()
59+
elif test == "train":
5860
train_args = split(
5961
f"--data {DATA_DIR}/coco128.data --img 416 --batch {self.batch_size} --nosave --notest \
6062
--epochs {self.num_epochs} --device {self.device_str} --weights '' \
@@ -64,8 +66,6 @@ def __init__(self, test, device, batch_size=None, extra_args=[]):
6466
self.training_loop, self.model, self.example_inputs = prepare_training_loop(
6567
train_args
6668
)
67-
elif test == "eval":
68-
self.model, self.example_inputs = self.prep_eval()
6969
self.amp_context = nullcontext
7070

7171
def prep_eval(self):

0 commit comments

Comments
 (0)