Fix TorchScript exporting for custom checkpoints

[zhiqwang]

An example to export the torchscript for the checkpoint trained from ultralytics/yolov5:

import cv2
import torch
from yolort.models import YOLO
from yolort.utils import read_image_to_tensor
from yolort.v5 import letterbox, attempt_download

# Loading and pre-processing the image
img_path = "test/assets/zidane.jpg"
img_raw = cv2.imread(img_path)
img = letterbox(img_raw, new_shape=(640, 640))[0]
img = read_image_to_tensor(img)

# will download yolov5n6.pt from
# <https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5n6.pt>
checkpoint_path = attempt_download("yolov5n6.pt")

model = YOLO.load_from_yolov5(checkpoint_path, score_thresh=0.25)
model.eval()
# Script the model
scripted_model = torch.jit.script(model)
scripted_model.eval()

out = model(img[None])
out_script = scripted_model(img[None])

torch.testing.assert_close(out[0]["scores"], out_script[1][0]["scores"], rtol=0, atol=0)
torch.testing.assert_close(out[0]["labels"], out_script[1][0]["labels"], rtol=0, atol=0)
torch.testing.assert_close(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0, atol=0)

cc @mattpopovich , this will fix #265 .

[CLAassistant]

All committers have signed the CLA.

[codecov]

Codecov Report

Merging #267 (87c1403) into main (9ebcea9) will increase coverage by 0.19%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #267      +/-   ##
==========================================
+ Coverage   93.62%   93.82%   +0.19%     
==========================================
  Files          11       11              
  Lines         675      696      +21     
==========================================
+ Hits          632      653      +21     
  Misses         43       43              

Flag	Coverage Δ
unittests	93.82% <100.00%> (+0.19%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
test/test_models.py	95.86% <100.00%> (+0.39%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9ebcea9...87c1403. Read the comment docs.