From 999463fbbdc23a9b4d944bdf74acba4e14f93d7c Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 30 Jan 2020 12:39:54 -0800
Subject: [PATCH] updates

---
 models.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/models.py b/models.py
index 4a57ef59..f77c1cdb 100755
--- a/models.py
+++ b/models.py
@@ -51,7 +51,11 @@ def create_modules(module_defs, img_size, arc):
                 modules = maxpool
 
         elif mdef['type'] == 'upsample':
-            modules = nn.Upsample(scale_factor=int(mdef['stride']), mode='nearest')
+            if ONNX_EXPORT:  # explicitly state size, avoid scale_factor
+                g = (yolo_index + 1) * 2
+                modules = nn.Upsample(size=(10 * g, 6 * g), mode='nearest')  # assume img_size = (320, 192)
+            else:
+                modules = nn.Upsample(scale_factor=int(mdef['stride']), mode='nearest')
 
         elif mdef['type'] == 'route':  # nn.Sequential() placeholder for 'route' layer
             layers = [int(x) for x in mdef['layers'].split(',')]