(self, minibatches, unlabeled=None)
| 702 | self.ema = self.hparams['mtl_ema'] |
| 703 | |
| 704 | def update(self, minibatches, unlabeled=None): |
| 705 | loss = 0 |
| 706 | for env, (x, y) in enumerate(minibatches): |
| 707 | loss += F.cross_entropy(self.predict(x, env), y) |
| 708 | |
| 709 | self.optimizer.zero_grad() |
| 710 | loss.backward() |
| 711 | self.optimizer.step() |
| 712 | |
| 713 | return {'loss': loss.item()} |
| 714 | |
| 715 | def update_embeddings_(self, features, env=None): |
| 716 | return_embedding = features.mean(0) |