Repository: Lightning-AI/torchmetrics

Issue:
nDCG can not be called with float targets
## 🐛 Bug

Invoking `RetrievalNormalizedNDCG` & `retrieval_normalized_dcg` with target of type float results into
`ValueError: `target` must be a tensor of booleans or integers`

The reason is this [check in `_check_retrieval_functional_inputs`](https://github.com/PyTorchLightning/metrics/blob/21fe0ca7e1e61e197a923c6482c5fa07e32908de/torchmetrics/utilities/checks.py#L514):
```py
if target.dtype not in (torch.bool, torch.long, torch.int):
    raise ValueError("`target` must be a tensor of booleans or integers")
```



### To Reproduce

The code samples below are sufficient to reproduce the error

#### Code sample

Code sample for `retrieval_normalized_dcg`
```py
import torch
from torchmetrics.functional import retrieval_normalized_dcg

preds = torch.tensor([.1, .2, .3, 4, 70])
target = torch.tensor([0.1, 0, 0, 0.5, 1.0])
retrieval_normalized_dcg(preds, target)
```

Code sample for `RetrievalNormalizedNDCG`
```py
import torch
from torchmetrics import RetrievalNormalizedDCG

indexes =torch.tensor([0, 0, 0, 1, 1, 1, 1])
preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
target = torch.tensor([0.1, 0.1, 0.5, 0.5, 1.0, 1.0, 0.0])
ndcg = RetrievalNormalizedDCG()
ndcg(preds, target, indexes=indexes)
```

### Expected behavior

Targets of type float are handled without errors

### Environment

- PyTorch Version (e.g., 1.0): 1.9.0
- OS (e.g., Linux): Ubuntu 20.04
- How you installed PyTorch (`conda`, `pip`, source): source
- Build command you used (if compiling from source): `pip install git+git://github.com/PyTorchLightning/metrics.git@79cb5e2f1744b0d4cd3f15adf9d725a68452a50`
- Python version: 3.8.10
- CUDA/cuDNN version: N/A
- GPU models and configuration: N/A
- Any other relevant information: N/A





Discussion:

Resolving patch:
diff --git a/torchmetrics/retrieval/retrieval_ndcg.py b/torchmetrics/retrieval/retrieval_ndcg.py
--- a/torchmetrics/retrieval/retrieval_ndcg.py
+++ b/torchmetrics/retrieval/retrieval_ndcg.py
@@ -27,7 +27,7 @@ class RetrievalNormalizedDCG(RetrievalMetric):
     Forward accepts:
 
     - ``preds`` (float tensor): ``(N, ...)``
-    - ``target`` (long or bool tensor): ``(N, ...)``
+    - ``target`` (long, int, bool or float tensor): ``(N, ...)``
     - ``indexes`` (long tensor): ``(N, ...)``
 
     ``indexes``, ``preds`` and ``target`` must have the same dimension.
diff --git a/torchmetrics/utilities/checks.py b/torchmetrics/utilities/checks.py
--- a/torchmetrics/utilities/checks.py
+++ b/torchmetrics/utilities/checks.py
@@ -503,7 +503,7 @@ def _check_retrieval_functional_inputs(
 
     Returns:
         preds: as torch.float32
-        target: as torch.long
+        target: as torch.long if not floating point else torch.float32
     """
     if preds.shape != target.shape:
         raise ValueError("`preds` and `target` must be of the same shape")
@@ -511,16 +511,7 @@ def _check_retrieval_functional_inputs(
     if not preds.numel() or not preds.size():
         raise ValueError("`preds` and `target` must be non-empty and non-scalar tensors")
 
-    if target.dtype not in (torch.bool, torch.long, torch.int):
-        raise ValueError("`target` must be a tensor of booleans or integers")
-
-    if not preds.is_floating_point():
-        raise ValueError("`preds` must be a tensor of floats")
-
-    if not allow_non_binary_target and (target.max() > 1 or target.min() < 0):
-        raise ValueError("`target` must contain `binary` values")
-
-    return preds.float().flatten(), target.long().flatten()
+    return _check_retrieval_target_and_prediction_types(preds, target, allow_non_binary_target=allow_non_binary_target)
 
 
 def _check_retrieval_inputs(
@@ -557,13 +548,38 @@ def _check_retrieval_inputs(
     if indexes.dtype is not torch.long:
         raise ValueError("`indexes` must be a tensor of long integers")
 
+    preds, target = _check_retrieval_target_and_prediction_types(
+        preds, target, allow_non_binary_target=allow_non_binary_target
+    )
+
+    return indexes.long().flatten(), preds, target
+
+
+def _check_retrieval_target_and_prediction_types(
+    preds: Tensor,
+    target: Tensor,
+    allow_non_binary_target: bool = False,
+) -> Tuple[Tensor, Tensor]:
+    """Check ``preds`` and ``target`` tensors are of the same shape and of the correct dtype.
+
+    Args:
+        preds: either tensor with scores/logits
+        target: tensor with ground true labels
+        allow_non_binary_target: whether to allow target to contain non-binary values
+
+    Raises:
+        ValueError:
+            If ``preds`` and ``target`` don't have the same shape, if they are empty
+            or not of the correct ``dtypes``.
+    """
+    if target.dtype not in (torch.bool, torch.long, torch.int) and not torch.is_floating_point(target):
+        raise ValueError("`target` must be a tensor of booleans, integers or floats")
+
     if not preds.is_floating_point():
         raise ValueError("`preds` must be a tensor of floats")
 
-    if target.dtype not in (torch.bool, torch.long, torch.int):
-        raise ValueError("`target` must be a tensor of booleans or integers")
-
     if not allow_non_binary_target and (target.max() > 1 or target.min() < 0):
         raise ValueError("`target` must contain `binary` values")
 
-    return indexes.long().flatten(), preds.float().flatten(), target.long().flatten()
+    target = target.float().flatten() if target.is_floating_point() else target.long().flatten()
+    return preds.float().flatten(), target


Test patch:
diff --git a/tests/retrieval/helpers.py b/tests/retrieval/helpers.py
--- a/tests/retrieval/helpers.py
+++ b/tests/retrieval/helpers.py
@@ -26,10 +26,11 @@
 from tests.retrieval.inputs import _input_retrieval_scores_all_target as _irs_all
 from tests.retrieval.inputs import _input_retrieval_scores_empty as _irs_empty
 from tests.retrieval.inputs import _input_retrieval_scores_extra as _irs_extra
+from tests.retrieval.inputs import _input_retrieval_scores_float_target as _irs_float_tgt
+from tests.retrieval.inputs import _input_retrieval_scores_int_target as _irs_int_tgt
 from tests.retrieval.inputs import _input_retrieval_scores_mismatching_sizes as _irs_mis_sz
 from tests.retrieval.inputs import _input_retrieval_scores_mismatching_sizes_func as _irs_mis_sz_fn
 from tests.retrieval.inputs import _input_retrieval_scores_no_target as _irs_no_tgt
-from tests.retrieval.inputs import _input_retrieval_scores_non_binary_target as _irs_non_binary
 from tests.retrieval.inputs import _input_retrieval_scores_wrong_targets as _irs_bad_tgt
 
 seed_all(42)
@@ -130,7 +131,6 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
         (_irs_empty.preds, _irs_empty.target, "`preds` and `target` must be non-empty and non-scalar tensors", {}),
         # check on input dtypes
         (_irs.preds.bool(), _irs.target, "`preds` must be a tensor of floats", {}),
-        (_irs.preds, _irs.target.float(), "`target` must be a tensor of booleans or integers", {}),
         # check targets are between 0 and 1
         (_irs_bad_tgt.preds, _irs_bad_tgt.target, "`target` must contain `binary` values", {}),
     ],
@@ -145,7 +145,6 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
         (_irs_empty.preds, _irs_empty.target, "`preds` and `target` must be non-empty and non-scalar tensors", {}),
         # check on input dtypes
         (_irs.preds.bool(), _irs.target, "`preds` must be a tensor of floats", {}),
-        (_irs.preds, _irs.target.float(), "`target` must be a tensor of booleans or integers", {}),
     ],
 )
 
@@ -228,13 +227,6 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
             "`preds` must be a tensor of floats",
             dict(empty_target_action="skip"),
         ),
-        (
-            _irs.indexes,
-            _irs.preds,
-            _irs.target.float(),
-            "`target` must be a tensor of booleans or integers",
-            dict(empty_target_action="skip"),
-        ),
     ],
 )
 
@@ -281,21 +273,6 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
             "`preds` must be a tensor of floats",
             dict(empty_target_action="skip"),
         ),
-        (
-            _irs.indexes,
-            _irs.preds,
-            _irs.target.float(),
-            "`target` must be a tensor of booleans or integers",
-            dict(empty_target_action="skip"),
-        ),
-        # check targets are between 0 and 1
-        (
-            _irs_bad_tgt.indexes,
-            _irs_bad_tgt.preds,
-            _irs_bad_tgt.target,
-            "`target` must contain `binary` values",
-            dict(empty_target_action="skip"),
-        ),
     ],
 )
 
@@ -321,7 +298,8 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
         (_irs.indexes, _irs.preds, _irs.target),
         (_irs_extra.indexes, _irs_extra.preds, _irs_extra.target),
         (_irs_no_tgt.indexes, _irs_no_tgt.preds, _irs_no_tgt.target),
-        (_irs_non_binary.indexes, _irs_non_binary.preds, _irs_non_binary.target),
+        (_irs_int_tgt.indexes, _irs_int_tgt.preds, _irs_int_tgt.target),
+        (_irs_float_tgt.indexes, _irs_float_tgt.preds, _irs_float_tgt.target),
     ],
 )
 
@@ -340,7 +318,8 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
         (_irs.preds, _irs.target),
         (_irs_extra.preds, _irs_extra.target),
         (_irs_no_tgt.preds, _irs_no_tgt.target),
-        (_irs_non_binary.preds, _irs_non_binary.target),
+        (_irs_int_tgt.preds, _irs_int_tgt.target),
+        (_irs_float_tgt.preds, _irs_float_tgt.target),
     ],
 )
 
diff --git a/tests/retrieval/inputs.py b/tests/retrieval/inputs.py
--- a/tests/retrieval/inputs.py
+++ b/tests/retrieval/inputs.py
@@ -32,12 +32,18 @@
     target=torch.randint(high=2, size=(NUM_BATCHES, BATCH_SIZE, EXTRA_DIM)),
 )
 
-_input_retrieval_scores_non_binary_target = Input(
+_input_retrieval_scores_int_target = Input(
     indexes=torch.randint(high=10, size=(NUM_BATCHES, 2 * BATCH_SIZE)),
     preds=torch.rand(NUM_BATCHES, 2 * BATCH_SIZE),
     target=torch.randint(low=-1, high=4, size=(NUM_BATCHES, 2 * BATCH_SIZE)),
 )
 
+_input_retrieval_scores_float_target = Input(
+    indexes=torch.randint(high=10, size=(NUM_BATCHES, 2 * BATCH_SIZE)),
+    preds=torch.rand(NUM_BATCHES, 2 * BATCH_SIZE),
+    target=torch.rand(NUM_BATCHES, 2 * BATCH_SIZE),
+)
+
 # with errors
 _input_retrieval_scores_no_target = Input(
     indexes=torch.randint(high=10, size=(NUM_BATCHES, BATCH_SIZE)),
diff --git a/tests/retrieval/test_ndcg.py b/tests/retrieval/test_ndcg.py
--- a/tests/retrieval/test_ndcg.py
+++ b/tests/retrieval/test_ndcg.py
@@ -117,6 +117,9 @@ def test_precision_gpu(self, indexes: Tensor, preds: Tensor, target: Tensor):
     def test_arguments_class_metric(
         self, indexes: Tensor, preds: Tensor, target: Tensor, message: str, metric_args: dict
     ):
+        if target.is_floating_point():
+            pytest.skip("NDCG metric works with float target input")
+
         self.run_metric_class_arguments_test(
             indexes=indexes,
             preds=preds,
@@ -135,6 +138,9 @@ def test_arguments_class_metric(
         )
     )
     def test_arguments_functional_metric(self, preds: Tensor, target: Tensor, message: str, metric_args: dict):
+        if target.is_floating_point():
+            pytest.skip("NDCG metric works with float target input")
+
         self.run_functional_metric_arguments_test(
             preds=preds,
             target=target,