lower the number of transformer layers in af3 during testing

Author: lucidrains

tests/test_af3.py

@@ -269,7 +269,24 @@ def test_alphafold3():
         dim_atom_inputs = 77,
         dim_additional_residue_feats = 33,
         dim_template_feats = 44,
-        num_dist_bins = 38
+        num_dist_bins = 38,
+        confidence_head_kwargs = dict(
+            pairformer_depth = 1
+        ),
+        template_embedder_kwargs = dict(
+            pairformer_stack_depth = 1
+        ),
+        msa_module_kwargs = dict(
+            depth = 1
+        ),
+        pairformer_stack = dict(
+            depth = 2
+        ),
+        diffusion_module_kwargs = dict(
+            atom_encoder_depth = 1,
+            token_transformer_depth = 1,
+            atom_decoder_depth = 1,
+        ),
     )
 
     loss = alphafold3(