Update stid_arch.py

Fix bugs in Issue #7 .
zezhishao · Jun 12, 2023 · f9801a5 · f9801a5
1 parent fef4ebc
commit f9801a5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stid/stid_arch/stid_arch.py b/stid/stid_arch/stid_arch.py
@@ -51,8 +51,8 @@ def __init__(self, **model_args):
 
         # encoding
         self.hidden_dim = self.embed_dim+self.node_dim * \
-            int(self.if_spatial)+self.temp_dim_tid*int(self.if_day_in_week) + \
-            self.temp_dim_diw*int(self.if_time_in_day)
+            int(self.if_spatial)+self.temp_dim_tid*int(self.if_time_in_day) + \
+            self.temp_dim_diw*int(self.if_day_in_week)
         self.encoder = nn.Sequential(
             *[MultiLayerPerceptron(self.hidden_dim, self.hidden_dim) for _ in range(self.num_layer)])