fix: 🐛 a bug in STAEformer dow embeddings (isse #219)

GestaltCogTeam · Jan 1, 2025 · 4b5f0ad · 4b5f0ad
1 parent d91d272
commit 4b5f0ad
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/baselines/STAEformer/arch/staeformer_arch.py b/baselines/STAEformer/arch/staeformer_arch.py
@@ -201,16 +201,16 @@ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_s
         batch_size = x.shape[0]
 
         if self.tod_embedding_dim > 0:
-            tod = x[..., 1]
+            tod = x[..., 1] * self.steps_per_day
         if self.dow_embedding_dim > 0:
-            dow = x[..., 2]
+            dow = x[..., 2] * 7
         x = x[..., : self.input_dim]
 
         x = self.input_proj(x)  # (batch_size, in_steps, num_nodes, input_embedding_dim)
         features = [x]
         if self.tod_embedding_dim > 0:
             tod_emb = self.tod_embedding(
-                (tod * self.steps_per_day).long()
+                tod.long()
             )  # (batch_size, in_steps, num_nodes, tod_embedding_dim)
             features.append(tod_emb)
         if self.dow_embedding_dim > 0: