def swap_state(model, state):
    """
    Context manager that swaps the state of a model, e.g:

        # model is in old state
        with swap_state(model, new_state):
            # model in new state
        # model back to old state
    """
    old_state = copy_state(model.state_dict())
    model.load_state_dict(state, strict=False)
    try:
        yield
    finally:
        model.load_state_dict(old_state)
Context manager that swaps the state of a model, e.g:

# model is in old state
with swap_state(model, new_state):
# model in new state
# model back to old state
demucs/states.py