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'