RoIHeads: remove low-scoring boxes statically

One of the postprocessing step of RoIHeads is to filter low-scoring (in
terms of confidence level) boxes from the output. This commit makes that
operation static for TPUs.

For now, it's needed to make it static manually as the experimental
support for dynamic nonzero/masked_select would raise the following
issue in this case:

RuntimeError: Internal: From /job:tpu_worker/replica:0/task:0:
RET_CHECK failure (learning/brain/google/xla/tpu_execute.cc:1066) size >= 0
	 [[{{node XRTExecute}}]]

torchvision/models/detection/roi_heads.py

@@ -595,8 +595,13 @@ def postprocess_detections(self, class_logits, box_regression, proposals, image_
             labels = labels.reshape(-1)
 
             # remove low scoring boxes
-            inds = torch.nonzero(scores > self.score_thresh).squeeze(1)
-            boxes, scores, labels = boxes[inds], scores[inds], labels[inds]
+            score_mask = (scores > self.score_thresh)
+            score_mask_boxes = score_mask.reshape(boxes.shape[0], -1)
+            nan_boxes = torch.full((1, boxes.shape[1],), float('nan'), device=device)
+            boxes = torch.where(score_mask_boxes, boxes, nan_boxes)
+            # pad with zeros so that nms will consider those boxes last
+            scores = torch.where(score_mask, scores, torch.tensor(0.0, device=device))
+            labels = torch.where(score_mask, labels, torch.tensor(-1, device=device))
 
             # remove empty boxes
             keep = box_ops.remove_small_boxes(boxes, min_size=1e-2)