Merge pull request meituan#801 from triple-Mu/triplemu/fix-upsample

Fix Upsample bug in infer..
ZHEQIUSHUI · Apr 28, 2023 · 62064fc · 62064fc
2 parents 041f547 + 7812aa1
commit 62064fc
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/yolov6/core/inferer.py b/yolov6/core/inferer.py
@@ -62,6 +62,8 @@ def model_switch(self, model, img_size):
         for layer in model.modules():
             if isinstance(layer, RepVGGBlock):
                 layer.switch_to_deploy()
+            elif isinstance(layer, torch.nn.Upsample) and not hasattr(layer, 'recompute_scale_factor'):
+                layer.recompute_scale_factor = None  # torch 1.11.0 compatibility
 
         LOGGER.info("Switch model to deploy modality.")