address #17

lucidrains · Nov 27, 2023 · 3a3f1e1 · 3a3f1e1
1 parent 20074ed
commit 3a3f1e1
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 5 deletions.
diff --git a/iTransformer/iTransformer.py b/iTransformer/iTransformer.py
@@ -207,7 +207,7 @@ def forward(
             assert self.training
             mse_loss = 0.
             for target, pred in zip(targets, pred_list):
-                assert targets.shape == pred_list.shape
+                assert target.shape == pred.shape
 
                 mse_loss = mse_loss + F.mse_loss(target, pred)
 

diff --git a/iTransformer/iTransformer2D.py b/iTransformer/iTransformer2D.py
@@ -313,7 +313,7 @@ def forward(
             assert self.training
             mse_loss = 0.
             for target, pred in zip(targets, pred_list):
-                assert targets.shape == pred_list.shape
+                assert target.shape == pred.shape
 
                 mse_loss = mse_loss + F.mse_loss(target, pred)
 

diff --git a/iTransformer/iTransformerFFT.py b/iTransformer/iTransformerFFT.py
@@ -230,7 +230,7 @@ def forward(
             assert self.training
             mse_loss = 0.
             for target, pred in zip(targets, pred_list):
-                assert targets.shape == pred_list.shape
+                assert target.shape == pred.shape
 
                 mse_loss = mse_loss + F.mse_loss(target, pred)
 

diff --git a/iTransformer/iTransformerNormConditioned.py b/iTransformer/iTransformerNormConditioned.py
@@ -243,7 +243,7 @@ def forward(
             assert self.training
             mse_loss = 0.
             for target, pred in zip(targets, pred_list):
-                assert targets.shape == pred_list.shape
+                assert target.shape == pred.shape
 
                 mse_loss = mse_loss + F.mse_loss(target, pred)
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'iTransformer',
   packages = find_packages(exclude=[]),
-  version = '0.4.1',
+  version = '0.4.2',
   license='MIT',
   description = 'iTransformer - Inverted Transformer Are Effective for Time Series Forecasting',
   author = 'Phil Wang',