Remove non-trainable modules (LPIPS, tracker) from state dict before saving.
(self, **kwargs)
| 47 | return self |
| 48 | |
| 49 | def state_dict(self, **kwargs): |
| 50 | """Remove non-trainable modules (LPIPS, tracker) from state dict before saving.""" |
| 51 | state_dict = super().state_dict(**kwargs) |
| 52 | for k in list(state_dict.keys()): |
| 53 | if 'lpips_loss' in k or 'tracker_prior' in k: |
| 54 | del state_dict[k] |
| 55 | return state_dict |
| 56 | |
| 57 | def load_state_dict(self, state_dict, strict=True): |
| 58 | """Load state dict, filtering out non-trainable modules if present.""" |