code: def _fix_post_backward_hook():
      global _hook_fixed
      if _hook_fixed:
          return
      _hook_fixed = True
  
      from torch.distributed.fsdp import _runtime_utils
      from torch.distributed.fsdp._common_utils import (
          TrainingState,
          HandleTrainingState,
      )
  
      old_hook = _runtime_utils._post_backward_hook
  
      def _post_backward_hook(state, handle, *args, **kwargs):
          checkpointed = getattr(
              state._fsdp_wrapped_module,
              "_audiocraft_checkpointed",
              False,
          )
          if checkpointed:
              # there will be one more forward in the backward with checkpointing and that will
              # massively confuse FSDP, so we have to make it think everything
              # is going according to the plan.
              state.training_state = TrainingState.FORWARD_BACKWARD
              handle._training_state = HandleTrainingState.BACKWARD_PRE
          old_hook(state, handle, *args, **kwargs)
  
      _runtime_utils._post_backward_hook = _post_backward_hook
  
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 6 for calling a global function with example usage: _fix_post_backward_hook() without return types