address #20

iTransformer/iTransformerNormConditioned.py

@@ -121,6 +121,7 @@ def __init__(
         super().__init__()
         self.num_variates = num_variates
         self.lookback_len = lookback_len
+        self.num_tokens_per_variate = num_tokens_per_variate
 
         self.mem_tokens = nn.Parameter(torch.randn(num_mem_tokens, dim)) if num_mem_tokens > 0 else None
 
@@ -228,8 +229,12 @@ def forward(
 
         # denormalize
 
+        x = rearrange(x, 'b (v n) d -> n b v d', n = self.num_tokens_per_variate)
+
         x = (x * var.sqrt()) + mean
 
+        x = rearrange(x, 'n b v d -> b (v n) d')
+
         # predicting multiple times
 
         pred_list = [fn(x) for fn in self.pred_heads]

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'iTransformer',
   packages = find_packages(exclude=[]),
-  version = '0.5.1',
+  version = '0.5.2',
   license='MIT',
   description = 'iTransformer - Inverted Transformer Are Effective for Time Series Forecasting',
   author = 'Phil Wang',