Skip to content

Commit a603cdd

Browse files
fmassafacebook-github-bot
authored andcommitted
[fbsync] Add missing device info. (#3651)
Reviewed By: NicolasHug Differential Revision: D27706953 fbshipit-source-id: c3810183d8acb7feec6868e28705ee25806acc27
1 parent 5af2cf0 commit a603cdd

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

torchvision/models/detection/retinanet.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -386,7 +386,8 @@ def compute_loss(self, targets, head_outputs, anchors):
386386
matched_idxs = []
387387
for anchors_per_image, targets_per_image in zip(anchors, targets):
388388
if targets_per_image['boxes'].numel() == 0:
389-
matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64))
389+
matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64,
390+
device=anchors_per_image.device))
390391
continue
391392

392393
match_quality_matrix = box_ops.box_iou(targets_per_image['boxes'], anchors_per_image)

0 commit comments

Comments
 (0)