fix a bug, addressing #6 thanks to @MattMcPartlon

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -194,23 +194,19 @@ def forward(
         time = None
     ):
         h = self.heads
-        has_context = exists(context)
 
-        context = default(context, x)
-
-        if x.shape[-1] != self.norm.gamma.shape[-1]:
-            print(context.shape, x.shape, self.norm.gamma.shape)
+        if exists(context):
+            context = self.norm_context(context)
 
         x = self.norm(x)
 
+        context = default(context, x)
+
         if exists(self.time_cond):
             assert exists(time)
             scale, shift = self.time_cond(time).chunk(2, dim = -1)
             x = (x * (scale + 1)) + shift
 
-        if has_context:
-            context = self.norm_context(context)
-
         qkv = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
 

setup.py

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