(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, use_checkpoint=False)
| 319 | p_drop=p_drop) |
| 320 | |
| 321 | def forward(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, use_checkpoint=False): |
| 322 | rbf_feat = rbf(torch.cdist(xyz[:,:,1,:], xyz[:,:,1,:])) |
| 323 | if use_checkpoint: |
| 324 | msa = checkpoint.checkpoint(create_custom_forward(self.msa2msa), msa, pair, rbf_feat, state) |
| 325 | pair = checkpoint.checkpoint(create_custom_forward(self.msa2pair), msa, pair) |
| 326 | pair = checkpoint.checkpoint(create_custom_forward(self.pair2pair), pair, rbf_feat) |
| 327 | R, T, state, alpha = checkpoint.checkpoint(create_custom_forward(self.str2str, top_k=0), msa, pair, R_in, T_in, xyz, state, idx, motif_mask) |
| 328 | else: |
| 329 | msa = self.msa2msa(msa, pair, rbf_feat, state) |
| 330 | pair = self.msa2pair(msa, pair) |
| 331 | pair = self.pair2pair(pair, rbf_feat) |
| 332 | R, T, state, alpha = self.str2str(msa, pair, R_in, T_in, xyz, state, idx, motif_mask=motif_mask, top_k=0) |
| 333 | |
| 334 | return msa, pair, R, T, state, alpha |
| 335 | |
| 336 | class IterativeSimulator(nn.Module): |
| 337 | def __init__(self, n_extra_block=4, n_main_block=12, n_ref_block=4, |
nothing calls this directly
no test coverage detected