From 7309a1ed7963904dddb633a07766422fa1e69f7c Mon Sep 17 00:00:00 2001 From: kevin Date: Tue, 2 Mar 2021 17:47:36 +0800 Subject: [PATCH] add torchscript --- src/models/backbones/mobilenetv2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/backbones/mobilenetv2.py b/src/models/backbones/mobilenetv2.py index 2b72db7..709d352 100644 --- a/src/models/backbones/mobilenetv2.py +++ b/src/models/backbones/mobilenetv2.py @@ -136,7 +136,7 @@ class MobileNetV2(nn.Module): # Initialize weights self._init_weights() - def forward(self, x, feature_names=None): + def forward(self, x): # Stage1 x = self.features[0](x) x = self.features[1](x)