code: def autocast_fn(
      enabled: bool, amp_precision: Optional[str] = "float16"
  ):
      if enabled:
          # If AMP is enabled, ensure that:
          # 1. Device is CUDA
          # 2. dtype is FLOAT16 or BFLOAT16
          if amp_precision not in str_to_torch_dtype:
              logger.error(
                  "For Mixed-precision training, supported dtypes are {}. Got: {}".format(
                      list(str_to_torch_dtype.keys()), amp_precision
                  )
              )
  
          if not torch.cuda.is_available():
              logger.error(
                  "For mixed-precision training, CUDA device is required."
              )
  
          return autocast(
              enabled=enabled, dtype=str_to_torch_dtype[amp_precision]
          )
      else:
          return autocast(enabled=False)
  
file_path: ml-cvnets/engine/utils.py
imports: import argparse
  from typing import Dict, List, Optional, Union
  import torch
  from torch import Tensor
  from torch.cuda.amp import autocast
  from utils import logger
  from utils.common_utils import create_directories
  from utils.ddp_utils import is_master
label_desc: Use image node_id 1 for calling a global function with example usage: autocast_fn(enabled, amp_precision) and returns: autocast, autocast