Update tal.py

mct-2.1.1
Laughing-q 3 months ago
parent 0d45fd8ef4
commit 4ded1eca39
  1. 2
      ultralytics/utils/tal.py

@ -306,7 +306,7 @@ def make_anchors(feats, strides, grid_cell_offset=0.5):
assert feats is not None
dtype, device = feats[0].dtype, feats[0].device
for i, stride in enumerate(strides):
h, w = feats[i].shape[2:] if isinstance(feats, list) else (int(feats[i]), int(feats[i]))
h, w = feats[i].shape[2:] if isinstance(feats, list) else (int(feats[i][0]), int(feats[i][1]))
sx = torch.arange(end=w, device=device, dtype=dtype) + grid_cell_offset # shift x
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)

Loading…
Cancel
Save