Fix `TRAIN_ON_PRED_BOXES` setting
Summary: Detaching predicted boxes from gradient computation since D2 assumes that `proposal_boxes.tensor` has `require_grad=False`
@@ -633,6 +633,7 @@ class StandardROIHeads(ROIHeads):
)
if self.training:
if self.train_on_pred_boxes:
+ with torch.no_grad():
pred_boxes = outputs.predict_boxes_for_gt_classes()
for proposals_per_image, pred_boxes_per_image in zip(proposals, pred_boxes):
proposals_per_image.proposal_boxes = Boxes(pred_boxes_per_image)
