Skip to content

Commit 186d207

Browse files
committed
[fix] bbox_pred
1 parent 2b88d5d commit 186d207

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

rtdetr_pytorch/src/zoo/rtdetr/rtdetr_postprocessor.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,12 @@ def forward(self, outputs, orig_target_sizes):
4747
else:
4848
scores = F.softmax(logits)[:, :, :-1]
4949
scores, labels = scores.max(dim=-1)
50+
boxes = bbox_pred
5051
if scores.shape[1] > self.num_top_queries:
5152
scores, index = torch.topk(scores, self.num_top_queries, dim=-1)
5253
labels = torch.gather(labels, dim=1, index=index)
5354
boxes = torch.gather(boxes, dim=1, index=index.unsqueeze(-1).tile(1, 1, boxes.shape[-1]))
54-
55+
5556
# TODO for onnx export
5657
if self.deploy_mode:
5758
return labels, boxes, scores

0 commit comments

Comments
 (0)