Skip to content

Commit

Permalink
Ensemble() visualize fix (ultralytics#3973)
Browse files Browse the repository at this point in the history
* fix visualize error

* Revert "fix visualize error"

* add visualise profile
  • Loading branch information
seven320 authored Jul 11, 2021
1 parent 5b8992f commit 4b9f76f
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions models/experimental.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4b9f76f

Please sign in to comment.