code: def switch_to_full_state_dict(models: tp.List[FSDP]):
      # Another bug in FSDP makes it that we cannot use the `state_dict_type` API,
      # so let's do thing manually.
      for model in models:
          FSDP.set_state_dict_type(  # type: ignore
              model,
              StateDictType.FULL_STATE_DICT,
              FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
          )
      try:
          yield
      finally:
          for model in models:
              FSDP.set_state_dict_type(
                  model, StateDictType.LOCAL_STATE_DICT
              )
  
file_path: audiocraft/audiocraft/optim/fsdp.py
imports: from contextlib import contextmanager
  import typing
  import dora
  import torch
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, FullStateDictConfig, StateDictType
  from torch.distributed._shard.sharded_tensor.api import ShardedTensor
label_desc: Use image node_id 3 for calling a global function with example usage: switch_to_full_state_dict(models) without return types