From fbccd3076d2f6b654d2e4adb2ecb7e6e49970f80 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Wed, 30 Mar 2022 21:49:26 +0200 Subject: [PATCH] expand static feat --- tft/module.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tft/module.py b/tft/module.py index 55afb67..b570432 100644 --- a/tft/module.py +++ b/tft/module.py @@ -225,10 +225,11 @@ class TemporalFusionDecoder(nn.Module): def forward( self, x: torch.Tensor, static: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> torch.Tensor: - static = static.repeat((1, self.context_length + self.prediction_length, 1)) + expanded_static = static.expand_as(x) + # static.repeat((1, self.context_length + self.prediction_length, 1)) skip = x[:, self.context_length :, ...] - x = self.enrich(x, static) + x = self.enrich(x, expanded_static) # does not work on GPU :-( # mask_pad = torch.ones_like(mask)[:, 0:1, ...]