address #21

README.md

@@ -23,6 +23,7 @@ from phenaki_pytorch import CViViT
 cvivit = CViViT(
     dim = 512,
     codebook_size = 5000,
+    image_size = 256,
     patch_size = 32,
     temporal_patch_size = 2,
     spatial_depth = 4,

phenaki_pytorch/cvivit_trainer.py

@@ -73,8 +73,8 @@ def __init__(
         valid_frac = 0.05,
         random_split_seed = 42,
         ema_beta = 0.995,
-        ema_update_after_step = 500,
-        ema_update_every = 10,
+        ema_update_after_step = 0,
+        ema_update_every = 1,
         apply_grad_penalty_every = 4,
         accelerate_kwargs: dict = dict()
     ):
@@ -86,6 +86,7 @@ def __init__(
         self.vae = vae
 
         if self.is_main:
+            print('ema')
             self.ema_vae = EMA(vae, update_after_step = ema_update_after_step, update_every = ema_update_every)
 
         self.register_buffer('steps', torch.Tensor([0]))
@@ -250,6 +251,7 @@ def train_step(self):
         # update exponential moving averaged generator
 
         if self.is_main:
+            print('hmmmm')
             self.ema_vae.update()
 
         # sample results every so often

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'phenaki-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.38',
+  version = '0.0.39',
   license='MIT',
   description = 'Phenaki - Pytorch',
   author = 'Phil Wang',
@@ -21,7 +21,7 @@
     'accelerate',
     'beartype',
     'einops>=0.6',
-    'ema-pytorch',
+    'ema-pytorch>=0.1.0',
     'opencv-python',
     'pillow',
     'numpy',