code: def _trace(model: torch.nn.Module) -> torch.fx.GraphModule:
      """Traces the given model and automatically wraps untracable modules into leaves."""
      leaf_modules = set()
      tracer = _ExtendedLeafTracer(leaf_modules)
      for name, module in model.named_modules():
          # TODO(ehotaj): The default is_leaf_module includes everything in torch.nn.
          # This means that some coarse modules like nn.TransformerEncoder are treated
          # as leaves, not traced, and are unable to be sharded. We may want to extend our
          # sharding code to trace through these modules as well.
          if tracer.is_leaf_module(module, ""):
              continue
          try:
              tracer.trace(module)
          except (TypeError, torch.fx.proxy.TraceError):
              leaf_modules.add(module)
              tracer = _ExtendedLeafTracer(leaf_modules)
      graph = tracer.trace(model)
      return torch.fx.GraphModule(model, graph)
  
file_path: fairscale/fairscale/experimental/nn/auto_shard.py
imports: import logging
  from typing import Dict, List, Set
  import torch
  import torch.fx
  from torch.fx.node import Node
label_desc: Use image node_id 4 for calling a global function with example usage: _trace(model) and returns: torch