code: def setup_model():
      """Sets up a model for training or testing and log the results."""
      # Build the model
      model = builders.build_model()
  
      if cfg.FSDP.ENABLED:
          fsdp_config = {}
          fsdp_config[
              "reshard_after_forward"
          ] = cfg.FSDP.RESHARD_AFTER_FW
          fsdp_config["mixed_precision"] = cfg.TRAIN.MIXED_PRECISION
          if cfg.TRAIN.MIXED_PRECISION:
              fsdp_config["clear_autocast_cache"] = True
  
          # Enable LAYER_NORM_FP32 wrapping for mixed precision training only. It is not
          # required for full precision training.
          ema_model = builders.build_model()
          if cfg.TRAIN.MIXED_PRECISION and cfg.FSDP.LAYER_NORM_FP32:
  
              def do_wrap(m):
                  m = auto_wrap_ln(m)
                  return FSDP(m, **fsdp_config)
  
              model = do_wrap(model)
              ema_model = do_wrap(ema_model)
          else:
              with enable_wrap(wrapper_cls=FSDP, **fsdp_config):
                  model = FSDP(model, **fsdp_config)
                  ema_model = FSDP(ema_model, **fsdp_config)
  
      else:
          # Build the model
          model = builders.build_model()
          ema_model = deepcopy(model)
  
      logger.info("Model:\n{}".format(model)) if cfg.VERBOSE else ()
      # Log model complexity
      logger.info(
          logging.dump_log_data(net.complexity(model), "complexity")
      )
      # Transfer the model to the current GPU device
      cur_device = torch.cuda.current_device()
      model = model.cuda(device=cur_device)
      ema_model = ema_model.cuda(device=cur_device)
      # Use multi-process data parallel model in the multi-gpu setting
      if cfg.NUM_GPUS > 1 and not cfg.FSDP.ENABLED:
          # Make model replica operate on the current device
          ddp = torch.nn.parallel.DistributedDataParallel
          model = ddp(
              module=model,
              device_ids=[cur_device],
              output_device=cur_device,
          )
      return model, ema_model
  
imports: import contextlib
  import os
  import random
  from copy import deepcopy
  import numpy
  import pycls.core.benchmark
  import pycls.core.builders
  import pycls.core.checkpoint
  import pycls.core.config
  import pycls.core.distributed
  import pycls.core.logging
  import pycls.core.meters
  import pycls.core.net
  import pycls.core.optimizer
  import pycls.datasets.loader
  import torch
  import torch.cuda.amp
  from fairscale.nn import auto_wrap, config_auto_wrap_policy, enable_wrap
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.optim.grad_scaler import ShardedGradScaler
  from pycls.core.config import cfg
  from pycls.core.io import cache_url, pathmgr
label_desc: Use image node_id 3 for calling a global function with example usage: setup_model() and returns: model,  ema_model, FSDP