Skip to content

GitLab

Unverified Commit 07fb8ba7 authored by Vasilis Vryniotis's avatar Vasilis Vryniotis Committed by GitHub
Browse files

Add missing device info. (#3651)

parent ad9cc62a
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
torchvision/models/detection/retinanet.py
View file @ 07fb8ba7
...@@ -386,7 +386,8 @@ class RetinaNet(nn.Module): ...@@ -386,7 +386,8 @@ class RetinaNet(nn.Module):
matched_idxs = [] matched_idxs = []
for anchors_per_image, targets_per_image in zip(anchors, targets): for anchors_per_image, targets_per_image in zip(anchors, targets):
if targets_per_image['boxes'].numel() == 0: if targets_per_image['boxes'].numel() == 0:
matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64)) matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64,
device=anchors_per_image.device))
continue continue
match_quality_matrix = box_ops.box_iou(targets_per_image['boxes'], anchors_per_image) match_quality_matrix = box_ops.box_iou(targets_per_image['boxes'], anchors_per_image)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment