(self, context=None)
| 80 | |
| 81 | @contextmanager |
| 82 | def ema_scope(self, context=None): |
| 83 | if self.use_ema: |
| 84 | self.model_ema.store(self.parameters()) |
| 85 | self.model_ema.copy_to(self) |
| 86 | if context is not None: |
| 87 | logpy.info(f"{context}: Switched to EMA weights") |
| 88 | try: |
| 89 | yield None |
| 90 | finally: |
| 91 | if self.use_ema: |
| 92 | self.model_ema.restore(self.parameters()) |
| 93 | if context is not None: |
| 94 | logpy.info(f"{context}: Restored training weights") |
| 95 | |
| 96 | @abstractmethod |
| 97 | def encode(self, *args, **kwargs) -> torch.Tensor: |
no test coverage detected