i* needs gradients

Author: lucidrains

perfusion_pytorch/perfusion.py

@@ -312,7 +312,7 @@ def forward(
 
             # exponential moving average for concept input encoding
 
-            concept_text_enc = ema_concept_text_enc * decay + concept_text_enc * (1. - decay)
+            ema_concept_text_enc = ema_concept_text_enc * decay + concept_text_enc * (1. - decay)
 
             # store
 
@@ -321,7 +321,7 @@ def forward(
 
             # update ema i_*
 
-            self.ema_concept_text_encs[concept_id].data.copy_(concept_text_enc)
+            self.ema_concept_text_encs[concept_id].data.copy_(ema_concept_text_enc)
 
         else:
             assert self.initted[concept_id_tensor].all(), 'you have not initialized or trained this module for the concepts id given'
@@ -330,6 +330,12 @@ def forward(
 
         i, o, W = self.ema_concept_text_encs[concept_id_tensor], self.concept_outputs[concept_id_tensor], weights
 
+        # if training, i* needs gradients. use straight-through?
+        # check with author about this
+
+        if self.training:
+            i = i + concept_text_enc - concept_text_enc.detach()
+
         # main contribution eq (3)
 
         i_energy = opt_einsum('c o, o i, c i -> c', i, Ci, i)

setup.py

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