| 37 | """ |
| 38 | |
| 39 | def __init__( |
| 40 | self, |
| 41 | ema_decay: Union[None, float] = None, |
| 42 | monitor: Union[None, str] = None, |
| 43 | input_key: str = "jpg", |
| 44 | ): |
| 45 | super().__init__() |
| 46 | |
| 47 | self.input_key = input_key |
| 48 | self.use_ema = ema_decay is not None |
| 49 | if monitor is not None: |
| 50 | self.monitor = monitor |
| 51 | |
| 52 | if self.use_ema: |
| 53 | self.model_ema = LitEma(self, decay=ema_decay) |
| 54 | logpy.info(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") |
| 55 | |
| 56 | if version.parse(torch.__version__) >= version.parse("2.0.0"): |
| 57 | self.automatic_optimization = False |
| 58 | |
| 59 | # def apply_ckpt(self, ckpt: Union[None, str, dict]): |
| 60 | # if ckpt is None: |