From 647223a7a8e5d8fca69481d477fe9a2930a8004c Mon Sep 17 00:00:00 2001
From: KEN <33506506+seven320@users.noreply.github.com>
Date: Mon, 12 Jul 2021 02:47:08 +0900
Subject: [PATCH] `Ensemble()` visualize fix (#3973)

* fix visualize error

* Revert "fix visualize error"

* add visualise profile
---
 models/experimental.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/models/experimental.py b/models/experimental.py
index d316b18..30dc361 100644
--- 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