code: def optimize_model(
      model,
      device,
      local_rank,
      optimizer=None,
      distributed=False,
      use_amp=None,
  ):
      """
      Wraps MultiGPU or distributed usage around a model
      No support for ONNX models
  
      :param model: model to optimize (e.g. trimming weights to fp16 / mixed precision)
      :type model: AdaptiveModel
      :param device: either gpu or cpu, get the device from initialize_device_settings()
      :param distributed: Whether training on distributed machines
      :param local_rank: rank of the machine in a distributed setting
      :param use_amp: Optimization level of nvidia's automatic mixed precision (AMP). The higher the level, the faster the model.
                      Options:
                      "O0" (Normal FP32 training)
                      "O1" (Mixed Precision => Recommended)
                      "O2" (Almost FP16)
                      "O3" (Pure FP16).
                      See details on: https://nvidia.github.io/apex/amp.html
      :return: model, optimizer
      """
      model, optimizer = _init_amp(model, device, optimizer, use_amp)
  
      if distributed:
          if APEX_PARALLEL_AVAILABLE:
              model = convert_syncbn_model(model)
              logger.info(
                  "Multi-GPU Training via DistributedDataParallel and apex.parallel"
              )
          else:
              logger.info(
                  "Multi-GPU Training via DistributedDataParallel"
              )
  
          # for some models DistributedDataParallel might complain about parameters
          # not contributing to loss. find_used_parameters remedies that.
          model = WrappedDDP(
              model,
              device_ids=[local_rank],
              output_device=local_rank,
              find_unused_parameters=True,
          )
  
      elif torch.cuda.device_count() > 1 and device.type == "cuda":
          model = WrappedDataParallel(model)
          logger.info("Multi-GPU Training via DataParallel")
  
      return model, optimizer
  
file_path: FARM/farm/modeling/optimization.py
imports: from importlib import import_module
  import logging
  import sys
  import inspect
  import torch
  from torch.nn.parallel import DistributedDataParallel
  from torch.nn import DataParallel
  from transformers.optimization import AdamW
  from farm.utils import MLFlowLogger
label_desc: Use image node_id 4 for calling a global function with example usage: optimize_model(model, device, local_rank, optimizer, distributed, use_amp) and returns: model,  optimizer