HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.1.28 avoid * ops on bool Tensors for RT-DETR OpenVINO export (#8937)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Kayzwer 2024-03-15 08:20:10 +08:00 • committed by GitHub
parent bc34ae560b
commit 2d513a9e4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 2 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/nn/modules/head.py
View file

@ -395,7 +395,7 @@ class RTDETRDecoder(nn.Module):
anchors.append(torch.cat([grid_xy, wh], -1).view(-1, h * w, 4)) # (1, h*w, 4)
anchors = torch.cat(anchors, 1) # (1, h*w*nl, 4)
valid_mask = ((anchors > eps) * (anchors < 1 - eps)).all(-1, keepdim=True) # 1, h*w*nl, 1
valid_mask = ((anchors > eps) & (anchors < 1 - eps)).all(-1, keepdim=True) # 1, h*w*nl, 1
anchors = torch.log(anchors / (1 - anchors))
anchors = anchors.masked_fill(~valid_mask, float("inf"))
return anchors, valid_mask