From a9933dc4fff88528ace8d3c51d6db54049cda6de Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Wed, 30 Mar 2022 17:45:23 +0200 Subject: [PATCH] fix selection --- tft/module.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tft/module.py b/tft/module.py index 4ed3219..4fa76c5 100644 --- a/tft/module.py +++ b/tft/module.py @@ -232,9 +232,8 @@ class TemporalFusionDecoder(nn.Module): # key_padding_mask = torch.cat((mask, mask_pad), dim=1).bool() query_key_value = x - attn_output, _ = self.attention( - query=query_key_value[-self.prediction_length :, ...], + query=query_key_value[:, -self.prediction_length :, ...], key=query_key_value, value=query_key_value, # key_padding_mask=key_padding_mask,