fix issue with ema class

README.md

@@ -499,10 +499,12 @@ loss.backward()
 
 ### DALL-E2 with Latent Diffusion
 
-This repository decides to take the next step and offer DALL-E2 combined with <a href="https://huggingface.co/spaces/multimodalart/latentdiffusion">latent diffusion</a>, from Rombach et al.
+This repository decides to take the next step and offer DALL-E v2 combined with <a href="https://huggingface.co/spaces/multimodalart/latentdiffusion">latent diffusion</a>, from Rombach et al.
 
 You can use it as follows. Latent diffusion can be limited to just the first U-Net in the cascade, or to any number you wish.
 
+The repository also comes equipped with all the necessary settings to recreate `ViT-VQGan` from the <a href="https://arxiv.org/abs/2110.04627">Improved VQGans</a> paper. Furthermore, the <a href="https://github.com/lucidrains/vector-quantize-pytorch">vector quantization</a> library also comes equipped to do <a href="https://arxiv.org/abs/2203.01941">residual or multi-headed quantization</a>, which I believe will give an even further boost in performance to the autoencoder.
+
 ```python
 import torch
 from dalle2_pytorch import Unet, Decoder, CLIP, VQGanVAE

dalle2_pytorch/train.py

@@ -35,7 +35,7 @@ class EMA(nn.Module):
 
         self.update_moving_average(self.ema_model, self.online_model)
 
-    def update_moving_average(ma_model, current_model):
+    def update_moving_average(self, ma_model, current_model):
         def calculate_ema(beta, old, new):
             if not exists(old):
                 return new

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.0.55',
+  version = '0.0.56',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',