diff --git a/ovi/modules/model.py b/ovi/modules/model.py
index f1c16a5..3f94ded 100644
--- a/ovi/modules/model.py
+++ b/ovi/modules/model.py
@@ -689,6 +689,7 @@ def set_rope_params(self):
     def set_gradient_checkpointing(self, enable: bool):
         self.gradient_checkpointing = enable
 
+    @torch.compiler.disable
     def prepare_transformer_block_kwargs(
         self,
         x,
@@ -802,7 +803,8 @@ def prepare_transformer_block_kwargs(
             context_lens=context_lens)
 
         return x, e, kwargs
-        
+
+    @torch.compiler.disable
     def post_transformer_block_out(self, x, grid_sizes, e):
         # head
         x = self.head(x, e)