minor fix

2022-07-21 20:37:15 +03:00
parent 894a93e318
commit 2596994f39
2 changed files with 2 additions and 2 deletions
@@ -59,7 +59,7 @@ class Detect(nn.Module):
                    y = torch.cat((xy, wh, y[..., 4:]), -1)
                z.append(y.view(bs, -1, self.no))

-        return x if self.training else torch.cat(z, 1)
+        return x if self.training else (torch.cat(z, 1), x)

    @staticmethod
    def _make_grid(nx=20, ny=20):