wassname/pytorch-transformer-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-transformer-ts.git synced 2026-07-21 12:50:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

expand static feat

Browse Source
This commit is contained in:
Kashif Rasul
2022-03-30 21:49:26 +02:00
parent 6e9f7af37b
commit fbccd3076d
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tft/module.py
+3 -2
View File
@@ -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, ...]