update: 适配昇腾

2025-11-27 09:57:34 +00:00 · 2025-11-27 09:57:34 +00:00 · 696c1b0793
commit 696c1b0793
parent e74b035c02
5 changed files with 34 additions and 11 deletions
--- a/ultralytics/utils/tal.py
+++ b/ultralytics/utils/tal.py
@ -341,7 +341,8 @@ def make_anchors(feats, strides, grid_cell_offset=0.5):
        sy = torch.arange(end=h, device=device, dtype=dtype) + grid_cell_offset  # shift y
        sy, sx = torch.meshgrid(sy, sx, indexing="ij") if TORCH_1_10 else torch.meshgrid(sy, sx)
        anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2))
-        stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))
+        # stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))
+        stride_tensor.append(torch.ones((h * w, 1), dtype=dtype, device=device)*stride)
    return torch.cat(anchor_points), torch.cat(stride_tensor)