From ff862d43cae80838c4ac72981d1a88f8322e601e Mon Sep 17 00:00:00 2001 From: Faakhir30 Date: Tue, 28 Jul 2026 13:16:56 +0500 Subject: [PATCH] fix samformer quantile Signed-off-by: Faakhir30 --- .../models/samformer/_samformer_v2.py | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/pytorch_forecasting/models/samformer/_samformer_v2.py b/pytorch_forecasting/models/samformer/_samformer_v2.py index 7d859921e..23f86d47b 100644 --- a/pytorch_forecasting/models/samformer/_samformer_v2.py +++ b/pytorch_forecasting/models/samformer/_samformer_v2.py @@ -97,9 +97,12 @@ def __init__( self.compute_values = nn.Linear( self.max_encoder_length, self.max_encoder_length ) # noqa: E501 + # Emit a separate forecast head per quantile so each has its own + # parameters and gradients (unlike tiling a point forecast with expand). self.linear_forecaster = nn.Linear( - self.max_encoder_length, self.max_prediction_length - ) # noqa: E501 + self.max_encoder_length, + self.max_prediction_length * self.n_quantiles, + ) def _scaled_dot_product_attention( self, @@ -169,18 +172,15 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: out = x_norm + att_score out = self.linear_forecaster(out) + # (batch, channels, prediction_length * n_quantiles) + batch_size, n_channels, _ = out.shape + out = out.view( + batch_size, + n_channels, + self.max_prediction_length, + self.n_quantiles, + ) + # Target is the last channel: (batch, prediction_length, n_quantiles) + target_predictions = out[:, -1, :, :] - out = out.transpose(1, 2) - - target_predictions = out[:, :, -1] # (batch_size, max_prediction_length) - - if target_predictions.ndim == 1: - target_predictions = target_predictions.unsqueeze(0) - - if self.n_quantiles > 1: - target_predictions = target_predictions.unsqueeze(-1).expand( - -1, -1, self.n_quantiles - ) - elif self.n_quantiles == 1: - target_predictions = target_predictions.unsqueeze(-1) return {"prediction": target_predictions}