[BUG] Safe fallback for retrieving index from QuantileLoss.to_prediction when 0.5 is unavailable (#2256)

References : Fixes #2255

## Description
If a user creates an instance of `QuantileLoss` defining bounding target
quantiles (e.g., `[0.1, 0.9]`) but excluding the direct median
probability `0.5`, performing a `.to_prediction()` pass crashes parsing
the output.
## Changes Made
- Introduced a conditional flow handling. It first checks `if 0.5 in
self.quantiles`, returning exact. If absent, it establishes the minimum
offset matching `0.5`, safely reverting to the closest approximated
center. Modified `pytorch_forecasting/metrics/quantile.py`.

CC : @phoeenniixx , @PranavBhatP , @fkiraly

---------

Co-authored-by: Aryan Saini <116151399+phoeenniixx@users.noreply.github.com>
Co-authored-by: Felix Hirwa Nshuti <hirwanshutiflx@gmail.com>
Co-authored-by: Franz Király <fkiraly@gcos.ai>

Author: 4 people

pytorch_forecasting/metrics/quantile.py

@@ -56,7 +56,13 @@ def to_prediction(self, y_pred: torch.Tensor) -> torch.Tensor:
             point prediction
         """
         if y_pred.ndim == 3:
-            idx = self.quantiles.index(0.5)
+            if 0.5 in self.quantiles:
+                idx = self.quantiles.index(0.5)
+            else:
+                idx = min(
+                    range(len(self.quantiles)),
+                    key=lambda i: abs(self.quantiles[i] - 0.5),
+                )
             y_pred = y_pred[..., idx]
         return y_pred
 

tests/test_metrics.py

@@ -19,6 +19,7 @@
     MultivariateNormalDistributionLoss,
     NegativeBinomialDistributionLoss,
     NormalDistributionLoss,
+    QuantileLoss,
 )
 from pytorch_forecasting.metrics.base_metrics import (
     AggregationMetric,
@@ -570,3 +571,22 @@ def test_MASE():
 
     assert scaling.shape == (batch_size,)
     assert (scaling > 0).all(), "Scaling should be positive"
+
+
+def test_QuantileLoss_to_prediction_fallback():
+    """Test to_prediction selects median when present, nearest quantile otherwise."""
+
+    loss_with_median = QuantileLoss(quantiles=[0.1, 0.5, 0.9])
+    y_pred_3d = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
+    result = loss_with_median.to_prediction(y_pred_3d)
+    expected = torch.tensor([[2.0, 5.0]])
+    assert torch.equal(result, expected)
+
+    loss_no_median = QuantileLoss(quantiles=[0.1, 0.4, 0.9])
+    result_fallback = loss_no_median.to_prediction(y_pred_3d)
+    expected_fallback = torch.tensor([[2.0, 5.0]])
+    assert torch.equal(result_fallback, expected_fallback)
+
+    y_pred_2d = torch.tensor([[10.0, 20.0]])
+    result_2d = loss_no_median.to_prediction(y_pred_2d)
+    assert torch.equal(result_2d, y_pred_2d)