Added python model tests for retinanet 50 (#2803)

ebe3f023 · vfdev · GitHub · 2831f11a · ebe3f023 · ebe3f023
Unverified Commit ebe3f023 authored Oct 14, 2020 by vfdev Committed by GitHub Oct 14, 2020
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

test/test_models.py test/test_models.py +3 -0

torchvision/models/detection/retinanet.py torchvision/models/detection/retinanet.py +1 -1

No files found.
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -71,6 +71,9 @@ script_test_models = {
    "keypointrcnn_resnet50_fpn": {
        'unwrapper': lambda x: x[1]
    },
+    "retinanet_resnet50_fpn": {
+        'unwrapper': lambda x: x[1]
+    }
 }

--- a/torchvision/models/detection/retinanet.py
+++ b/torchvision/models/detection/retinanet.py
@@ -565,7 +565,7 @@ class RetinaNet(nn.Module):
            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)