code: def swap_state(model, state, **kwargs):
      old_state = copy_state(model.state_dict())
      model.load_state_dict(state, **kwargs)
      try:
          yield
      finally:
          model.load_state_dict(old_state)
  
imports: from concurrent.futures import ProcessPoolExecutor
  from contextlib import contextmanager
  from functools import wraps, lru_cache
  import hashlib
  import json
  import logging
  from pathlib import Path
  import typing
  import flashy
  import flashy.distrib
  import omegaconf
  import torch
  from torch.nn.utils.rnn import pad_sequence
label_desc: Use image node_id 15 for calling a global function with example usage: swap_state(model, state) without return types