Update head.py

8 months ago · 96e911c370
parent 159971ae59
commit 96e911c370
1 changed files with 1 additions and 0 deletions
--- a/ultralytics/nn/modules/head.py
+++ b/ultralytics/nn/modules/head.py
@ -121,6 +121,7 @@ class Detect(nn.Module):
        x2 = self.relu(dbox[:, 3, :])
        y_bb = torch.stack((x1, y1, x2, y2), 1).transpose(1, 2)

+        # return torch.cat((dbox, cls.sigmoid()), 1)
        return y_bb, cls.sigmoid().permute(0, 2, 1)

    def bias_init(self):