Commit 4c37c386 authored by Gustaf Ahdritz's avatar Gustaf Ahdritz
Browse files

Add conditions to model constructor

parent 6b547775
...@@ -79,24 +79,29 @@ class AlphaFold(nn.Module): ...@@ -79,24 +79,29 @@ class AlphaFold(nn.Module):
self.recycling_embedder = RecyclingEmbedder( self.recycling_embedder = RecyclingEmbedder(
**self.config["recycling_embedder"], **self.config["recycling_embedder"],
) )
self.template_angle_embedder = TemplateAngleEmbedder(
**self.template_config["template_angle_embedder"], if(self.template_config.enabled):
) self.template_angle_embedder = TemplateAngleEmbedder(
self.template_pair_embedder = TemplatePairEmbedder( **self.template_config["template_angle_embedder"],
**self.template_config["template_pair_embedder"], )
) self.template_pair_embedder = TemplatePairEmbedder(
self.template_pair_stack = TemplatePairStack( **self.template_config["template_pair_embedder"],
**self.template_config["template_pair_stack"], )
) self.template_pair_stack = TemplatePairStack(
self.template_pointwise_att = TemplatePointwiseAttention( **self.template_config["template_pair_stack"],
**self.template_config["template_pointwise_attention"], )
) self.template_pointwise_att = TemplatePointwiseAttention(
self.extra_msa_embedder = ExtraMSAEmbedder( **self.template_config["template_pointwise_attention"],
**self.extra_msa_config["extra_msa_embedder"], )
)
self.extra_msa_stack = ExtraMSAStack( if(self.extra_msa_config.enabled):
**self.extra_msa_config["extra_msa_stack"], self.extra_msa_embedder = ExtraMSAEmbedder(
) **self.extra_msa_config["extra_msa_embedder"],
)
self.extra_msa_stack = ExtraMSAStack(
**self.extra_msa_config["extra_msa_stack"],
)
self.evoformer = EvoformerStack( self.evoformer = EvoformerStack(
**self.config["evoformer_stack"], **self.config["evoformer_stack"],
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment