Update yolo.py

[pedrombmachado]

Replaced by the original yolo.py from Ultralytics. This solves the on-going problem of training the model.
replace the current forward Detect method by:

def forward(self, x):
        """
        Forward pass of YOLOv5 Detect head.
        - Training: reshapes output for loss computation
        - Export: returns raw conv outputs (no view/permute, for DPU compliance)
        - Inference: standard YOLOv5 decode with anchors/grids
        """
        z = []          # inference outputs
        outputs = []    # store outputs for all detection layers

        for i in range(self.nl):
            out = self.m[i](x[i])  # conv -> [bs, na*no, h, w]
            bs, _, ny, nx = out.shape

            if self.training:
                # reshape for loss
                out = out.view(bs, self.na, self.no, ny, nx) \
                        .permute(0, 1, 3, 4, 2) \
                        .contiguous()  # [bs, na, ny, nx, no]
                outputs.append(out)

            elif self.export:
                # Export for Vitis-AI: raw conv outputs only
                outputs.append(out)  # [bs, na*no, h, w]

            else:
                # Normal inference path
                out = out.view(bs, self.na, self.no, ny, nx) \
                        .permute(0, 1, 3, 4, 2) \
                        .contiguous()  # [bs, na, ny, nx, no]

                if not self.training:  # same as YOLOv5
                    if self.dynamic or self.grid[i].shape[2:4] != out.shape[2:4]:
                        self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

                    # split outputs into xy, wh, conf
                    xy, wh, conf = out.sigmoid().split((2, 2, self.nc + 1), 4)
                    xy = (xy * 2 + self.grid[i]) * self.stride[i]     # xy coords
                    wh = (wh * 2) ** 2 * self.anchor_grid[i]          # wh coords
                    y = torch.cat((xy, wh, conf), 4)

                    z.append(y.view(bs, self.na * nx * ny, self.no))

                outputs.append(out)

        if self.training:
            return outputs                           # list of [bs, na, ny, nx, no]
        elif self.export:
            return outputs                           # list of [bs, na*no, h, w]
        else:
            return (torch.cat(z, 1), outputs)        # inference return