AttributeError: 'AlphaFold' object has no attribute 'template_pair_embedder' when using long_sequence_inference=True

[Lix1993]

Default Model fails with OOM, when using long_sequence_inference=True， it breaks at template offload

---

AttributeError Traceback (most recent call last) Cell In[13], line 2 1 t = time.perf_counter() ----> 2 out = model(processed_feature_dict) 3 print(f"Inference time: {time.perf_counter() - t}")

File ~/workspace/envs/af2_openfold_official/conda_env/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs) 1126 # If we don't have any hooks, we want to skip the rest of the logic in 1127 # this function, and just call forward. 1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks 1129 or _global_forward_hooks or _global_forward_pre_hooks): -> 1130 return forward_call(*input, **kwargs) 1131 # Do not call functions when jit is used 1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/workspace/envs/af2_openfold_official/openfold/openfold/model/model.py:568, in AlphaFold.forward(self, batch) 565 torch.clear_autocast_cache() 567 # Run the next iteration of the model --> 568 outputs, m_1_prev, z_prev, x_prev, early_stop = self.iteration( 569 feats, 570 prevs, 571 _recycle=(num_iters > 1) 572 ) 574 num_recycles += 1 576 if not is_final_iter:

File ~/workspace/envs/af2_openfold_official/openfold/openfold/model/model.py:325, in AlphaFold.iteration(self, feats, prevs, recycle) 320 if self.config.template.enabled: 321 template_feats = { 322 k: v for k, v in feats.items() if k.startswith("template") 323 } --> 325 template_embeds = self.embed_templates( 326 template_feats, 327 feats, 328 z, 329 pair_mask.to(dtype=z.dtype), 330 no_batch_dims, 331 inplace_safe=inplace_safe, 332 ) 334 # [*, N, N, C_z] 335 z = add(z, 336 template_embeds.pop("template_pair_embedding"), 337 inplace_safe, 338 )

File ~/workspace/envs/af2_openfold_official/openfold/openfold/model/model.py:159, in AlphaFold.embed_templates(self, batch, feats, z, pair_mask, templ_dim, inplace_safe) 157 else: 158 if self.template_config.offload_templates: --> 159 return embed_templates_offload(self, 160 batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe, 161 ) 162 elif self.template_config.average_templates: 163 return embed_templates_average(self, 164 batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe, 165 )

File ~/workspace/envs/af2_openfold_official/openfold/openfold/model/template.py:523, in embed_templates_offload(model, batch, z, pair_mask, templ_dim, template_chunk_size, inplace_safe) 515 # [, N, N, C_t] 516 t = build_template_pair_feat( 517 single_template_feats, 518 use_unit_vector=model.config.template.use_unit_vector, (...) 521 **model.config.template.distogram, 522 ).to(z.dtype) --> 523 t = model.template_pair_embedder(t) 525 # [, 1, N, N, C_z] 526 t = model.template_pair_stack( 527 t.unsqueeze(templ_dim), 528 pair_mask.unsqueeze(-3).to(dtype=z.dtype), (...) 533 _mask_trans=model.config._mask_trans, 534 )

File ~/workspace/envs/af2_openfold_official/conda_env/lib/python3.9/site-packages/torch/nn/modules/module.py:1207, in Module.getattr(self, name) 1205 if name in modules: 1206 return modules[name] -> 1207 raise AttributeError("'{}' object has no attribute '{}'".format( 1208 type(self).name, name))

AttributeError: 'AlphaFold' object has no attribute 'template_pair_embedder'