Fix postprocessing in reference implementation (mlcommons#1164)

GATEOverflow · Jun 21, 2022 · 799a410 · 799a410
1 parent b284212
commit 799a410
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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]]