Repository: Lightning-AI/torchmetrics

Issue:
`RuntimeError`s when attempting use of `MAP`-metric
## 🐛 Bug

When using `MAP`-metric during `test_step` to update values and `test_epoch_end` to compute the values, the `MAP.compute` fails with `RuntimeError: Tensors must be CUDA and dense` in `all_gather`.

### To Reproduce

Pick a regional neural network, from for example lightning bolts' `FasterRCNN`. Implement `test_step` where 
`MAP.update()` is called for the model results and targets and the `MAP.compute()` is called in the `test_epoch_end` for the previously gathered results.

#### Code sample

```python
from pl_bolts.models.detection.faster_rcnn import FasterRCNN as _FasterRCNN
from torchmetrics.detection.map import MAP

class FasterRCNN(_FasterRCNN):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._voc_map = MAP(iou_thresholds=[0.3, 0.5, 0.7], class_metrics=True)
    # <I have here validation_step code, but this should not matter>
    def test_step(self, batch, batch_idx) -> Dict[str, float]:
        """Test function to ran over during test interference."""
        imgs, targets = batch
        outs = self.model(imgs)

        # If move to CPU is not done, error is `torch.sort does not support boolean tensor`.
        for target in targets:
            target["boxes"] = target["boxes"].cpu()
            target["labels"] = target["labels"].cpu()
        for target in outs:
            target["boxes"] = target["boxes"].cpu()
            target["labels"] = target["labels"].cpu()
            target["scores"] = target["scores"].cpu()
        # Update MAP metric created in `__init__`
        self._voc_map.update(
            outs,
            targets
        )
        # No reason to return anything as the values have been collected to `MAP`
        return None

    def test_epoch_end(self, output_results):
        """Test function ran after all test steps. Aggregates all test results to mAP."""        
        voc = self._voc_map.compute() # Fails with `RuntimeError: Tensors must be CUDA and dense`
        self.log("mAP", voc)
```
### Expected behavior

The `MAP.update` updates values correctly whether or not they had been transferred to GPU device and `compute` does not fail with `RuntimeError: Tensors must be CUDA and dense` during call to `torch.distributed.gather_all`.

### Environment

- PyTorch Version (e.g., 1.0): 1.9.1
- PL Bolts Version: 0.4.0
- TorchMetrics: 0.6.1
- OS (e.g., Linux): Ubuntu 18.04
- How you installed PyTorch (`conda`, `pip`, source): `poetry` installation
- Python version: 3.9.6
- CUDA: 11.3:
- NVIDIA GeForce 2080 Ti x 2

### Additional context

`distributed_backend=nccl`



Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @vahvero,
Could you try running master version of torchmetrics:
```pip install git+https://github.com/PytorchLightning/metrics.git@master```
We recently had a PR that fixed cuda issues for the MAP metric (https://github.com/PyTorchLightning/metrics/pull/665) so this issue may already be solved.

Resolving patch:
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -278,8 +278,7 @@ def setup(app):
 
 # copy all notebooks to local folder
 path_nbs = os.path.join(_PATH_HERE, "notebooks")
-if not os.path.isdir(path_nbs):
-    os.mkdir(path_nbs)
+os.makedirs(path_nbs, exist_ok=True)
 for path_ipynb in glob.glob(os.path.join(_PATH_ROOT, "notebooks", "*.ipynb")):
     path_ipynb2 = os.path.join(path_nbs, os.path.basename(path_ipynb))
     shutil.copy(path_ipynb, path_ipynb2)
diff --git a/torchmetrics/detection/map.py b/torchmetrics/detection/map.py
--- a/torchmetrics/detection/map.py
+++ b/torchmetrics/detection/map.py
@@ -423,10 +423,10 @@ def _evaluate_image(
         nb_iou_thrs = len(self.iou_thresholds)
         nb_gt = len(gt)
         nb_det = len(det)
-        gt_matches = torch.zeros((nb_iou_thrs, nb_gt), dtype=torch.bool)
-        det_matches = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool)
+        gt_matches = torch.zeros((nb_iou_thrs, nb_gt), dtype=torch.bool, device=self.device)
+        det_matches = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool, device=self.device)
         gt_ignore = ignore_area_sorted
-        det_ignore = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool)
+        det_ignore = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool, device=self.device)
         if torch.numel(ious) > 0:
             for idx_iou, t in enumerate(self.iou_thresholds):
                 for idx_det in range(nb_det):
@@ -436,7 +436,7 @@ def _evaluate_image(
                         det_matches[idx_iou, idx_det] = True
                         gt_matches[idx_iou, m] = True
         # set unmatched detections outside of area range to ignore
-        det_areas = box_area(det)
+        det_areas = box_area(det).to(self.device)
         det_ignore_area = (det_areas < area_range[0]) | (det_areas > area_range[1])
         ar = det_ignore_area.reshape((1, nb_det))
         det_ignore = torch.logical_or(
@@ -565,6 +565,10 @@ def _calculate(self, class_ids: List) -> Tuple[Dict, MAPMetricResults, MARMetric
         recall = -torch.ones((nb_iou_thrs, nb_classes, nb_bbox_areas, nb_max_det_thrs))
         scores = -torch.ones((nb_iou_thrs, nb_rec_thrs, nb_classes, nb_bbox_areas, nb_max_det_thrs))
 
+        # move tensors if necessary
+        self.max_detection_thresholds = self.max_detection_thresholds.to(self.device)
+        self.rec_thresholds = self.rec_thresholds.to(self.device)
+
         # retrieve E at each category, area range, and max number of detections
         for idx_cls in range(nb_classes):
             for idx_bbox_area in range(nb_bbox_areas):


Test patch:
diff --git a/tests/helpers/testers.py b/tests/helpers/testers.py
--- a/tests/helpers/testers.py
+++ b/tests/helpers/testers.py
@@ -25,6 +25,7 @@
 
 from torchmetrics import Metric
 from torchmetrics.detection.map import MAPMetricResults
+from torchmetrics.utilities.data import apply_to_collection
 
 try:
     set_start_method("spawn")
@@ -164,10 +165,8 @@ def _class_test(
 
     # move to device
     metric = metric.to(device)
-
-    if isinstance(preds, torch.Tensor):
-        preds = preds.to(device)
-        target = target.to(device)
+    preds = apply_to_collection(preds, Tensor, lambda x: x.to(device))
+    target = apply_to_collection(target, Tensor, lambda x: x.to(device))
 
     kwargs_update = {k: v.to(device) if isinstance(v, Tensor) else v for k, v in kwargs_update.items()}