remove groupnorms https://arxiv.org/abs/2312.02696

lucidrains · May 3, 2024 · 2f57546 · 2f57546
1 parent 16561b3
commit 2f57546
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/make_a_video_pytorch/make_a_video.py b/make_a_video_pytorch/make_a_video.py
@@ -346,12 +346,11 @@ def __init__(
         dim,
         dim_out,
         kernel_size = 3,
-        temporal_kernel_size = None,
-        groups = 8
+        temporal_kernel_size = None
     ):
         super().__init__()
         self.project = PseudoConv3d(dim, dim_out, 3)
-        self.norm = nn.GroupNorm(groups, dim_out)
+        self.norm = RMSNorm(dim_out)
         self.act = nn.SiLU()
 
     def forward(
@@ -376,7 +375,6 @@ def __init__(
         dim_out,
         *,
         timestep_cond_dim = None,
-        groups = 8
     ):
         super().__init__()
 
@@ -388,8 +386,8 @@ def __init__(
                 nn.Linear(timestep_cond_dim, dim_out * 2)
             )
 
-        self.block1 = Block(dim, dim_out, groups = groups)
-        self.block2 = Block(dim_out, dim_out, groups = groups)
+        self.block1 = Block(dim, dim_out)
+        self.block2 = Block(dim_out, dim_out)
         self.res_conv = PseudoConv3d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
 
     def forward(

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'make-a-video-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.3.1',
+  version = '0.4.0',
   license='MIT',
   description = 'Make-A-Video - Pytorch',
   author = 'Phil Wang',