(self, **kwargs)
| 41 | self.lpips_loss.requires_grad_(False) |
| 42 | |
| 43 | def state_dict(self, **kwargs): |
| 44 | # remove lpips_loss |
| 45 | state_dict = super().state_dict(**kwargs) |
| 46 | for k in list(state_dict.keys()): |
| 47 | if 'lpips_loss' in k: |
| 48 | del state_dict[k] |
| 49 | return state_dict |
| 50 | |
| 51 | def load_state_dict(self, state_dict, strict=True): |
| 52 | # Optionally, if the LPIPS keys exist in the state_dict, remove them before loading. |