code: def _create_model(
      with_model2,
      with_sync_bn,
      with_fsdp,
      with_checkpoint,
      mixed_precision,
      flatten,
      wrap_bn,
      fp32_reduce_scatter,
      bucket_cap_mb,
  ):
      model = Model2() if with_model2 else Model()
      fsdp_config = None
      if with_sync_bn:
          model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
          fsdp_config = {
              "mixed_precision": False,
              "flatten_parameters": False,
              "reshard_after_forward": False,
              "bucket_cap_mb": 0,
              "force_input_to_fp32": True,  # SyncBN needs this.
          }
  
      if with_fsdp and wrap_bn:
          model.block1 = auto_wrap_bn(
              model.block1,
              single_rank_pg=False,
              fsdp_config=fsdp_config,
          )
          model.block2 = auto_wrap_bn(
              model.block2,
              single_rank_pg=False,
              fsdp_config=fsdp_config,
          )
          if with_model2:
              model.block3 = auto_wrap_bn(
                  model.block3,
                  single_rank_pg=False,
                  fsdp_config=fsdp_config,
              )
  
      if with_checkpoint:
          model.block2 = checkpoint_wrapper(model.block2)
          if with_model2:
              model.block3 = checkpoint_wrapper(model.block3)
  
      if with_fsdp:
          with enable_wrap(
              wrapper_cls=FSDP,
              flatten_parameters=flatten,
              mixed_precision=mixed_precision,
              compute_dtype=torch.float32,
              fp32_reduce_scatter=fp32_reduce_scatter,
              bucket_cap_mb=bucket_cap_mb,
          ):
              model.block1 = wrap(model.block1)
              model.block2 = wrap(model.block2)
              if with_model2:
                  model.block3 = wrap(model.block3)
              model.head = wrap(model.head)
  
      return model
  
imports: import contextlib
  import pickle
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.nn import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import auto_wrap_bn
  from fairscale.nn.wrap import enable_wrap, wrap
label_desc: Use image node_id 1 for calling a global function with example usage: _create_model(with_model2, with_sync_bn, with_fsdp, with_checkpoint, mixed_precision, flatten, wrap_bn, fp32_reduce_scatter, bucket_cap_mb) and returns: model