Fix `Grad strides do not match bucket view strides` warning for YOLO12 DDP training (#19395)

2 weeks ago · c15fc952e1
parent 77eac7898a
commit c15fc952e1
1 changed files with 2 additions and 2 deletions
--- a/ultralytics/nn/modules/block.py
+++ b/ultralytics/nn/modules/block.py
@ -1227,8 +1227,8 @@ class AAttn(nn.Module):
            v = v.reshape(B // self.area, N * self.area, C)
            B, N, _ = x.shape

-        x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
-        v = v.reshape(B, H, W, C).permute(0, 3, 1, 2)
+        x = x.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()
+        v = v.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()

        x = x + self.pe(v)
        return self.proj(x)