From 0226dd9971c604e99614e98378edfa44acb1a756 Mon Sep 17 00:00:00 2001
From: Zhiqiang Wang <zhiqwang@outlook.com>
Date: Fri, 10 Sep 2021 03:15:24 -0400
Subject: [PATCH] Move test_torchscript to test_models

---
 test/test_models.py      | 21 ++++++++++++++++++++-
 test/test_torchscript.py | 23 -----------------------
 2 files changed, 20 insertions(+), 24 deletions(-)
 delete mode 100644 test/test_torchscript.py

diff --git a/test/test_models.py b/test/test_models.py
index b0da71e3..785ad7fe 100644
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -3,10 +3,11 @@
 import io
 import contextlib
 import warnings
-
+import pytest
 import torch
 from torch import Tensor
 
+from yolort import models
 from yolort.models.backbone_utils import darknet_pan_backbone
 from yolort.models.transformer import darknet_tan_backbone
 from yolort.models.anchor_utils import AnchorGenerator
@@ -257,3 +258,21 @@ def test_criterion(self):
         assert isinstance(losses['cls_logits'], Tensor)
         assert isinstance(losses['bbox_regression'], Tensor)
         assert isinstance(losses['objectness'], Tensor)
+
+
+@pytest.mark.parametrize('arch', ['yolov5s', 'yolov5m', 'yolov5l', 'yolotr'])
+def test_torchscript(arch):
+    model = models.__dict__[arch](pretrained=True, size=(320, 320), score_thresh=0.45)
+    model.eval()
+
+    scripted_model = torch.jit.script(model)
+    scripted_model.eval()
+
+    x = [torch.rand(3, 288, 320), torch.rand(3, 300, 256)]
+
+    out = model(x)
+    out_script = scripted_model(x)
+
+    torch.testing.assert_allclose(out[0]["scores"], out_script[1][0]["scores"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["labels"], out_script[1][0]["labels"], rtol=0., atol=0.)
+    torch.testing.assert_allclose(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0., atol=0.)
diff --git a/test/test_torchscript.py b/test/test_torchscript.py
deleted file mode 100644
index 69068ce8..00000000
--- a/test/test_torchscript.py
+++ /dev/null
@@ -1,23 +0,0 @@
-# Copyright (c) 2020, Zhiqiang Wang. All Rights Reserved.
-import pytest
-import torch
-
-from yolort import models
-
-
-@pytest.mark.parametrize('arch', ['yolov5s', 'yolov5m', 'yolov5l', 'yolotr'])
-def test_yolov5s_script(arch):
-    model = models.__dict__[arch](pretrained=True, size=(320, 320), score_thresh=0.45)
-    model.eval()
-
-    scripted_model = torch.jit.script(model)
-    scripted_model.eval()
-
-    x = [torch.rand(3, 288, 320), torch.rand(3, 300, 256)]
-
-    out = model(x)
-    out_script = scripted_model(x)
-
-    torch.testing.assert_allclose(out[0]["scores"], out_script[1][0]["scores"], rtol=0., atol=0.)
-    torch.testing.assert_allclose(out[0]["labels"], out_script[1][0]["labels"], rtol=0., atol=0.)
-    torch.testing.assert_allclose(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0., atol=0.)