def extract_states(sources: Architecture) -> Checkpoint:
targets = {}
for name, source in sources.items():
targets[name] = source.state_dict()
return targets
def load_states(sources: Checkpoint, targets: Architecture) -> None:
for name, target in targets.items():
target.load_state_dict(sources[name])