fix EMA updating buffers with non-float tensors

configs/train_decoder_config.example.json

@@ -15,7 +15,7 @@
         "channels": 3,
         "timesteps": 1000,
         "loss_type": "l2",
-        "beta_schedule": "cosine",
+        "beta_schedule": ["cosine"],
         "learned_variance": true
     },
     "data": {

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.11.3'
+__version__ = '0.11.4'

dalle2_pytorch/vqgan_vae_trainer.py

@@ -16,10 +16,11 @@ from torchvision.utils import make_grid, save_image
 
 from einops import rearrange
 
-from dalle2_pytorch.train import EMA
 from dalle2_pytorch.vqgan_vae import VQGanVAE
 from dalle2_pytorch.optimizer import get_optimizer
 
+from ema_pytorch import EMA
+
 # helpers
 
 def exists(val):
@@ -97,7 +98,7 @@ class VQGanVAETrainer(nn.Module):
         valid_frac = 0.05,
         random_split_seed = 42,
         ema_beta = 0.995,
-        ema_update_after_step = 2000,
+        ema_update_after_step = 500,
         ema_update_every = 10,
         apply_grad_penalty_every = 4,
         amp = False

setup.py

@@ -28,7 +28,7 @@ setup(
     'click',
     'clip-anytorch',
     'coca-pytorch>=0.0.5',
-    'ema-pytorch>=0.0.3',
+    'ema-pytorch>=0.0.7',
     'einops>=0.4',
     'einops-exts>=0.0.3',
     'embedding-reader',