Change reshape to support empty batches. (#3031)

1706921b · Vasilis Vryniotis · GitHub · 585ce2c4 · 1706921b
Unverified Commit 1706921b authored Nov 19, 2020 by Vasilis Vryniotis Committed by GitHub Nov 19, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

torchvision/models/detection/roi_heads.py torchvision/models/detection/roi_heads.py +1 -1

No files found.
--- a/torchvision/models/detection/roi_heads.py
+++ b/torchvision/models/detection/roi_heads.py
@@ -40,7 +40,7 @@ def fastrcnn_loss(class_logits, box_regression, labels, regression_targets):
    sampled_pos_inds_subset = torch.where(labels > 0)[0]
    labels_pos = labels[sampled_pos_inds_subset]
    N, num_classes = class_logits.shape
-    box_regression = box_regression.reshape(N, -1, 4)
+    box_regression = box_regression.reshape(N, box_regression.size(-1) // 4, 4)

    box_loss = det_utils.smooth_l1_loss(
        box_regression[sampled_pos_inds_subset, labels_pos],