Update YOLOv5 YAMLs to 'u' YAMLs (#800)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2023-02-04 19:54:34 +04:00
parent 0d182e80f1
commit 21ae321bc2
17 changed files with 68 additions and 46 deletions
--- a/ultralytics/nn/modules.py
+++ b/ultralytics/nn/modules.py
@ -456,4 +456,5 @@ class Classify(nn.Module):
    def forward(self, x):
        if isinstance(x, list):
            x = torch.cat(x, 1)
-        return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
+        x = self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
+        return x if self.training else x.softmax(1)