code: def fsdp_wrapper(module, **kwargs):
      """
      Customer FSDP wrapper, adding the missing options
      """
      from vissl.utils.layer_memory_tracking import ProcessGroupTracker
  
      # Add global process group to the list of keys
      fsdp_config = dict(**kwargs)
      if "process_group" not in fsdp_config:
          fsdp_config["process_group"] = get_global_group()
      if fsdp_config.get("_TRACK_COMMUNICATIONS", False):
          fsdp_config["process_group"] = ProcessGroupTracker(
              fsdp_config["process_group"]
          )
  
      # Remove keys that are not supported in FSDP
      for key in {
          "_TRACK_COMMUNICATIONS",
          "AUTO_WRAP_THRESHOLD",
          "FORCE_SYNC_CUDA",
      }:
          fsdp_config.pop(key, None)
  
      return FSDP(module, **fsdp_config)
  
file_path: vissl/vissl/utils/fsdp_utils.py
imports: import torch.distributed
  import torch.nn
  from fairscale.nn import auto_wrap, default_auto_wrap_policy, enable_wrap
  from fairscale.nn.data_parallel import auto_wrap_bn, FullyShardedDataParallel
  from vissl.config.attr_dict import AttrDict
label_desc: Use image node_id 6 for calling a global function with example usage: fsdp_wrapper(module) and returns: FSDP