added compilation to specific blocks

Cloud303 · Aug 21, 2024 · 5d8396b · 5d8396b
1 parent 9257dec
commit 5d8396b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/src/models/msa_module.py b/src/models/msa_module.py
@@ -229,6 +229,7 @@ def __init__(
             inf=inf
         )
 
+    @torch.compile
     def forward(
             self,
             m: Tensor,

diff --git a/src/models/pairformer.py b/src/models/pairformer.py
@@ -191,6 +191,7 @@ def __init__(
             transition_n,
         )
 
+    @torch.compile
     def forward(
             self,
             s: Tensor,  # (bs, 1, n_tokens, c_s)
-Original file line number
+Diff line change
@@ Expand Up / @@ -229,6 +229,7 @@ def __init__( @@
                 inf=inf
             )
+        @torch.compile
         def forward(
                 self,
                 m: Tensor,
@@ Expand Down @@