quick fix

lucidrains · Dec 1, 2023 · ceb9b9c · ceb9b9c
1 parent 907a320
commit ceb9b9c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'soundstorm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.3.0',
+  version = '0.3.1',
   license='MIT',
   description = 'SoundStorm - Efficient Parallel Audio Generation from Google Deepmind, in Pytorch',
   author = 'Phil Wang',

diff --git a/soundstorm_pytorch/soundstorm.py b/soundstorm_pytorch/soundstorm.py
@@ -221,7 +221,7 @@ def __init__(self, chan_in, chan_out, kernel_size, padding):
 
     def forward(self, x, mask = None):
         if exists(mask):
-            mask = mask[..., None]
+            mask = rearrange(mask, 'b n -> b 1 n')
             x = x.masked_fill(~mask, 0.)
 
         x = F.pad(x, self.padding)