Allow the executorch compilation.

kyutai-labs · Dec 20, 2024 · 5888f38 · 5888f38
1 parent 6aba1ae
commit 5888f38
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 4 deletions.
diff --git a/moshi/moshi/models/compression.py b/moshi/moshi/models/compression.py
@@ -143,7 +143,7 @@ def __init__(
         freeze_encoder: bool = False,
         freeze_quantizer: bool = False,
         freeze_quantizer_level: int = -1,
-        torch_compile_encoder_decoder: bool = False,
+        torch_compile_encoder_decoder: bool = True,
     ):
         super().__init__()
         self.encoder = encoder

diff --git a/moshi/moshi/modules/transformer.py b/moshi/moshi/modules/transformer.py
@@ -581,9 +581,6 @@ def _sa_block(self, x: torch.Tensor):
         return x_orig + self.layer_scale_1(update)
 
     def forward(self, x: torch.Tensor):
-        with ExitStack() as stack:
-            if x.device.type != 'cuda':
-                stack.enter_context(no_compile())
             x = self._sa_block(x)
             x = self._ff_block(x)
             state = self._streaming_state