fix sub-gn in mhesa

Author: lucidrains

README.md

@@ -73,14 +73,3 @@ logits = adapter(timeseries) # (1, 10)
     primaryClass = {cs.LG}
 }
 ```
-
-```bibtex
-@article{Sun2023RetentiveNA,
-    title   = {Retentive Network: A Successor to Transformer for Large Language Models},
-    author  = {Yutao Sun and Li Dong and Shaohan Huang and Shuming Ma and Yuqing Xia and Jilong Xue and Jianyong Wang and Furu Wei},
-    journal = {ArXiv},
-    year    = {2023},
-    volume. = {abs/2307.08621},
-    url     = {https://api.semanticscholar.org/CorpusID:259937453}
-}
-```

etsformer_pytorch/etsformer_pytorch.py

@@ -103,7 +103,11 @@ def __init__(
         self.dropout = nn.Dropout(dropout)
         self.alpha = nn.Parameter(torch.randn(heads))
 
-        self.norm_heads = nn.GroupNorm(heads, dim) if norm_heads else nn.Identity()
+        self.norm_heads = nn.Sequential(
+            Rearrange('b n (h d) -> b (h d) n', h = heads),
+            nn.GroupNorm(heads, dim),
+            Rearrange('b (h d) n -> b n (h d)', h = heads)
+        ) if norm_heads else nn.Identity()
 
         self.project_in = nn.Linear(dim, dim)
         self.project_out = nn.Linear(dim, dim)
@@ -178,10 +182,10 @@ def forward(self, x, naive = False):
 
         output = rearrange(output, 'b h n d -> b n (h d)')
 
-        # maybe groupnorm
-        # borrowing a trick from retnet paper
+        # maybe sub-ln from https://arxiv.org/abs/2210.06423 - retnet used groupnorm
 
         output = self.norm_heads(output)
+
         return self.project_out(output)
 
 ## frequency attention

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'ETSformer-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.0',
+  version = '0.1.1',
   license='MIT',
   description = 'ETSTransformer - Exponential Smoothing Transformer for Time-Series Forecasting - Pytorch',
   long_description_content_type = 'text/markdown',