add forward func

Wuhall · Dec 21, 2022 · 4dae398 · 4dae398
1 parent d8f6db5
commit 4dae398
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/model/backbone/MobileNetV3.py b/model/backbone/MobileNetV3.py
@@ -203,3 +203,8 @@ def forward_head(self, x, pre_logits: bool = False):
             if self.drop_rate > 0.:
                 x = F.dropout(x, p=self.drop_rate, training=self.training)
             return self.classifier(x)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.forward_head(x)
+        return x