add a final norm to the latents after all the latent self attention

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -278,6 +278,7 @@ def __init__(
         dim,
         latent_self_attn_depth,
         dim_latent = None,
+        final_norm = True,
         **attn_kwargs
     ):
         super().__init__()
@@ -292,6 +293,8 @@ def __init__(
                 FeedForward(dim_latent)
             ]))
 
+        self.latent_final_norm = LayerNorm(dim_latent) if final_norm else nn.Identity()
+
         self.patches_peg = PEG(dim)
         self.patches_self_attn = LinearAttention(dim, norm = True, **attn_kwargs)
         self.patches_self_attn_ff = FeedForward(dim)
@@ -323,6 +326,7 @@ def forward(self, patches, latents, t):
 
         patches = self.patches_cross_attn_ff(patches, time = t) + patches
 
+        latents = self.latent_final_norm(latents)
         return patches, latents
 
 class RIN(nn.Module):

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'RIN-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.0',
+  version = '0.1.1',
   license='MIT',
   description = 'RIN - Recurrent Interface Network - Pytorch',
   author = 'Phil Wang',