mirror of https://github.com/ZHKKKe/MODNet.git
add torchscript
parent
26033da8f6
commit
7309a1ed79
|
|
@ -136,7 +136,7 @@ class MobileNetV2(nn.Module):
|
||||||
# Initialize weights
|
# Initialize weights
|
||||||
self._init_weights()
|
self._init_weights()
|
||||||
|
|
||||||
def forward(self, x, feature_names=None):
|
def forward(self, x):
|
||||||
# Stage1
|
# Stage1
|
||||||
x = self.features[0](x)
|
x = self.features[0](x)
|
||||||
x = self.features[1](x)
|
x = self.features[1](x)
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue