change chunking threshold

Author: jwohlwend

Diff for: src/boltz/data/const.py

@@ -352,3 +352,10 @@
 
 max_msa_seqs = 16384
 max_paired_seqs = 8192
+
+
+####################################################################################################
+# CHUNKING
+####################################################################################################
+
+chunk_size_threshold = 384

Diff for: src/boltz/model/modules/trunk.py

@@ -224,7 +224,7 @@ def forward(
         """
         # Set chunk sizes
         if not self.training:
-            if z.shape[1] > 512:
+            if z.shape[1] > const.chunk_size_threshold:
                 chunk_heads_pwa = True
                 chunk_size_transition_z = 64
                 chunk_size_transition_msa = 32
@@ -520,7 +520,7 @@ def forward(
 
         """
         if not self.training:
-            if z.shape[1] > 512:
+            if z.shape[1] > const.chunk_size_threshold:
                 chunk_size_tri_attn = 128
             else:
                 chunk_size_tri_attn = 512