Ensemble() visualize fix (#3973)

* fix visualize error * Revert "fix visualize error" * add visualise profile
ultralytics · Jul 11, 2021 · 647223a · 647223a
1 parent a544d59
commit 647223a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/experimental.py b/models/experimental.py
@@ -100,10 +100,10 @@ class Ensemble(nn.ModuleList):
     def __init__(self):
         super(Ensemble, self).__init__()
 
-    def forward(self, x, augment=False):
+    def forward(self, x, augment=False, profile=False, visualize=False):
         y = []
         for module in self:
-            y.append(module(x, augment)[0])
+            y.append(module(x, augment, profile, visualize)[0])
         # y = torch.stack(y).max(0)[0]  # max ensemble
         # y = torch.stack(y).mean(0)  # mean ensemble
         y = torch.cat(y, 1)  # nms ensemble