Merge branch 'main' into add-nbeats-v2

Author: harshsomankar123-tech

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)