(self, state_dict, strict=True)
| 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. |
| 53 | filtered_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('lpips_loss')} |
| 54 | return super().load_state_dict(filtered_state_dict, strict=False) |
| 55 | |
| 56 | def train(self, mode=True): |
| 57 | super().train(mode) |