code: def purge_fsdp(model: FSDP):
      """Purge the FSDP cached shard inside the model. This should
      allow setting the best state or switching to the EMA.
      """
      from torch.distributed.fsdp._runtime_utils import _reshard  # type: ignore
  
      for module in FSDP.fsdp_modules(model):
          handles = module._handles
          if not handles:
              continue
          handle = handles[0]
          unsharded_flat_param = (
              handle._get_padded_unsharded_flat_param()
          )
          storage_size: int = unsharded_flat_param._typed_storage()._size()  # type: ignore
          if storage_size == 0:
              continue
          true_list = [True for h in handles]
          _reshard(module, handles, true_list)
  
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 5 for calling a global function with example usage: purge_fsdp(model) without return types