From 23d0456b08cac22f783d63292cc7c2bf87a19a60 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 10 Sep 2022 23:55:18 +0300
Subject: [PATCH] `torch.jit.trace()` fix (#9363)

* Update common.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

* Update ci-testing.yml

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---
 .github/workflows/ci-testing.yml | 3 +++
 models/common.py                 | 1 +
 2 files changed, 4 insertions(+)

diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml
index 6fb277676959..a83f997cbfc2 100644
--- a/.github/workflows/ci-testing.yml
+++ b/.github/workflows/ci-testing.yml
@@ -119,9 +119,12 @@ jobs:
           python export.py --weights $m.pt --img 64 --include torchscript  # export
           python - <<EOF
           import torch
+          im = torch.zeros([1, 3, 64, 64])
           for path in '$m', '$b':
               model = torch.hub.load('.', 'custom', path=path, source='local')
               print(model('data/images/bus.jpg'))
+              model(im)  # warmup, build grids for trace
+              torch.jit.trace(model, [im])
           EOF
       - name: Test classification
         shell: bash  # for Windows compatibility
diff --git a/models/common.py b/models/common.py
index c601aacc885c..8b7dbbfa95fe 100644
--- a/models/common.py
+++ b/models/common.py
@@ -600,6 +600,7 @@ def __init__(self, model, verbose=True):
         if self.pt:
             m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
             m.inplace = False  # Detect.inplace=False for safe multithread inference
+            m.export = True  # do not output loss values
 
     def _apply(self, fn):
         # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers