add torchscript

2021-03-02 17:47:36 +08:00 · 2021-03-02 17:47:36 +08:00 · 7309a1ed79
parent 26033da8f6
commit 7309a1ed79
1 changed files with 1 additions and 1 deletions
--- 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)