address #35

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'soundstorm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.4.3',
+  version = '0.4.6',
   license='MIT',
   description = 'SoundStorm - Efficient Parallel Audio Generation from Google Deepmind, in Pytorch',
   author = 'Phil Wang',

soundstorm_pytorch/soundstorm.py

@@ -530,6 +530,8 @@ def __init__(
             Rearrange('b n (h d) -> b (n h) d', h = num_effective_quantizers)
         )
 
+        self.num_effective_quantizers = num_effective_quantizers
+
         # each quantizer codebook would require its own logits weight and bias matrices
         # the amazing einops makes this easy with 'EinMix'
 
@@ -579,6 +581,9 @@ def forward(
         x = self.embedding_proj(x)
 
         if exists(sum_embeds):
+            if sum_embeds.ndim == 3:
+                sum_embeds = reduce(sum_embeds, 'b (n h) d -> b n d', 'sum', h = self.num_effective_quantizers)
+
             x = x + sum_embeds
 
         if exists(cond):