From 799a410e49d9390c34e63fab9bbe4a5ef03c3c48 Mon Sep 17 00:00:00 2001 From: Pablo Gonzalez Date: Tue, 21 Jun 2022 11:21:41 -0500 Subject: [PATCH] Fix postprocessing in reference implementation (#1164) --- vision/classification_and_detection/python/openimages.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vision/classification_and_detection/python/openimages.py b/vision/classification_and_detection/python/openimages.py index d1dc4ffe2..03fe0135a 100644 --- a/vision/classification_and_detection/python/openimages.py +++ b/vision/classification_and_detection/python/openimages.py @@ -246,9 +246,9 @@ def __call__(self, results, ids, expected=None, result_dict=None): if self.dict_format: # If the output of the model is in dictionary format. This happens # for the model ssd-resnext50-pytorch - bboxes_ = [e['boxes'] for e in results] - labels_ = [e['labels'] for e in results] - scores_ = [e['scores'] for e in results] + bboxes_ = [e['boxes'].cpu() for e in results] + labels_ = [e['labels'].cpu() for e in results] + scores_ = [e['scores'].cpu() for e in results] results = [bboxes_, labels_, scores_] else: bboxes_ = [results[0]]