allow for local attention to attend to nothing

audiolm_pytorch/soundstream.py

@@ -396,7 +396,16 @@ def __init__(
 
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                LocalMHA(dim = dim, heads = heads, qk_rmsnorm = True, window_size = window_size, use_rotary_pos_emb = not dynamic_pos_bias, use_xpos = True, **kwargs),
+                LocalMHA(
+                    dim = dim,
+                    heads = heads,
+                    qk_rmsnorm = True,
+                    window_size = window_size,
+                    use_rotary_pos_emb = not dynamic_pos_bias,
+                    gate_values_per_head = True,
+                    use_xpos = True,
+                    **kwargs
+                ),
                 FeedForward(dim = dim)
             ]))
 
@@ -610,7 +619,7 @@ def __init__(
         self.adversarial_loss_weight = adversarial_loss_weight
         self.feature_loss_weight = feature_loss_weight
 
-        self.register_buffer('zero', torch.tensor([0.]), persistent = False)
+        self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
     @property
     def device(self):

audiolm_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.6.2'
+__version__ = '1.6.3'

setup.py

@@ -27,7 +27,7 @@
     'fairseq',
     'joblib',
     'lion-pytorch',
-    'local-attention>=1.8.4',
+    'local-attention>=1.9.0',
     'scikit-learn',
     'sentencepiece',
     'torch>=1.12',