| 250 | [self.confidence_model.output_dim] * num_mlp_layer + [1]) |
| 251 | |
| 252 | def predict_rmsd(self, batch, all_loss=None, metric=None): |
| 253 | protein = batch['graph'] |
| 254 | if self.graph_construction_model: |
| 255 | protein = self.graph_construction_model(protein) |
| 256 | atom_feature = self.confidence_model(protein, protein.node_feature.float())["node_feature"] |
| 257 | residue_feature = scatter_mean(atom_feature, protein.atom2residue, dim=0, |
| 258 | dim_size=protein.num_residue) # [num_residue, feature_dim] |
| 259 | pred = self.mlp(residue_feature).squeeze(-1) # [num_residue] |
| 260 | return pred |
| 261 | |
| 262 | @torch.no_grad() |
| 263 | def generate(self, batch, randomize=True): |