Change batched NMS threshold to choose for-loop version (#4990)

According to the benchmark link https://github.com/pytorch/vision/issues/1311#issuecomment-781329339

, for GPU the threshold should be higher (which is why detectron2 used 40k).
This PR changes the threshold to be device dependent.
Co-authored-by: Nicolas Hug <contact@nicolas-hug.com>
Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

torchvision/ops/boxes.py

@@ -66,8 +66,7 @@ def batched_nms(
     _log_api_usage_once("torchvision.ops.batched_nms")
     # Benchmarks that drove the following thresholds are at
     # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339
-    # Ideally for GPU we'd use a higher threshold
-    if boxes.numel() > 4_000 and not torchvision._is_tracing():
+    if boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing():
         return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
     else:
         return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)