tweak

Author: lucidrains

perfusion_pytorch/perfusion.py

@@ -174,6 +174,13 @@ def forward(
                 # for the prompt ids not initialized yet, hard copy over the initial superclass outputs
                 self.superclass_output.data.copy_(superclass_output)
 
+            elif exists(superclass_output):
+                # if text enc with superclass is passed in for more than 1 batch
+                # just take the opportunity to exponentially average it a bit more
+
+                ema_superclass_output = self.superclass_output * decay + superclass_output * (1. - decay)
+                self.superclass_output.data.copy_(ema_superclass_output)
+
             # if any in the batch is not initialized, initialize
 
             if not initted:

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.17',
+  version = '0.0.18',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',