use latest gateloop

lucidrains · Dec 22, 2023 · c4fb6f6 · c4fb6f6
1 parent b379498
commit c4fb6f6
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 5 deletions.
diff --git a/audiolm_pytorch/soundstream.py b/audiolm_pytorch/soundstream.py
@@ -488,7 +488,6 @@ def __init__(
         attn_xpos_scale_base = None,
         attn_dynamic_pos_bias = False,
         use_gate_loop_layers = False,
-        gate_loop_use_jax = False,
         squeeze_excite = False,
         complex_stft_discr_logits_abs = True,
         pad_mode = 'reflect',
@@ -521,7 +520,7 @@ def __init__(
             encoder_blocks.append(EncoderBlock(chan_in, chan_out, layer_stride, enc_cycle_dilations, squeeze_excite, pad_mode))
 
             if use_gate_loop_layers:
-                encoder_blocks.append(Residual(ChannelTranspose(GateLoop(chan_out, use_jax_associative_scan = gate_loop_use_jax))))
+                encoder_blocks.append(Residual(ChannelTranspose(GateLoop(chan_out, use_heinsen = True))))
 
         self.encoder = nn.Sequential(
             CausalConv1d(input_channels, channels, 7, pad_mode = pad_mode),

diff --git a/audiolm_pytorch/version.py b/audiolm_pytorch/version.py
@@ -1 +1 @@
-__version__ = '2.0.2'
+__version__ = '2.0.3'
diff --git a/setup.py b/setup.py
@@ -26,13 +26,13 @@
     'encodec',
     'fairseq',
     'wandb',
-    'gateloop-transformer>=0.0.24',
+    'gateloop-transformer>=0.5.2',
     'joblib',
     'local-attention>=1.9.0',
     'pytorch-warmup',
     'scikit-learn',
     'sentencepiece',
-    'torch>=1.12',
+    'torch>=2.1',
     'torchaudio',
     'transformers',
     'tqdm',