(self, batch)
| 59 | self.train_chi_id = train_chi_id |
| 60 | |
| 61 | def forward(self, batch): |
| 62 | all_loss = torch.tensor(0, dtype=torch.float32, device=self.device) |
| 63 | metric = {} |
| 64 | |
| 65 | # Sample from the schedule |
| 66 | protein = batch['graph'] |
| 67 | t = self.schedule_1pi_periodic.sample_train_t(shape=(protein.batch_size,)).to(self.device) |
| 68 | |
| 69 | # Add noise to protein |
| 70 | train_chi_id = np.random.randint(self.NUM_CHI_ANGLES) if self.train_chi_id is None else self.train_chi_id |
| 71 | batch = self.add_noise(batch, t, chi_id=train_chi_id) |
| 72 | |
| 73 | # Predict and take loss |
| 74 | pred = self.predict(batch, all_loss, metric) |
| 75 | target = self.target(batch) |
| 76 | |
| 77 | metric = self.evaluate(pred, target) |
| 78 | all_loss += metric["diffusion loss"] |
| 79 | |
| 80 | return all_loss, metric |
| 81 | |
| 82 | def add_noise(self, batch, t, chi_id=None): |
| 83 | """Add noise to protein and update protein |
nothing calls this directly
no test coverage detected