Added python model tests for retinanet 50 (pytorch#2803)

bryant1410 · Nov 22, 2020 · e9be62d · e9be62d
1 parent 9768861
commit e9be62d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/test/test_models.py b/test/test_models.py
@@ -71,6 +71,9 @@ def get_available_video_models():
     "keypointrcnn_resnet50_fpn": {
         'unwrapper': lambda x: x[1]
     },
+    "retinanet_resnet50_fpn": {
+        'unwrapper': lambda x: x[1]
+    }
 }
 
 

diff --git a/torchvision/models/detection/retinanet.py b/torchvision/models/detection/retinanet.py
@@ -565,7 +565,7 @@ def forward(self, images, targets=None):
             if not self._has_warned:
                 warnings.warn("RetinaNet always returns a (Losses, Detections) tuple in scripting")
                 self._has_warned = True
-            return (losses, detections)
+            return losses, detections
         return self.eager_outputs(losses, detections)