code: def torch_dtype_from_precision(
      precision: Union[int, str], megatron_amp_O2: Optional[bool] = None
  ) -> torch.dtype:
      """Mapping from PTL precision types to corresponding PyTorch parameter datatype."""
      if megatron_amp_O2 is not None and megatron_amp_O2 is False:
          return torch.float32
  
      if precision in ["bf16", "bf16-mixed"]:
          return torch.bfloat16
      elif precision in [16, "16", "16-mixed"]:
          return torch.float16
      elif precision in [32, "32", "32-true"]:
          return torch.float32
      else:
          raise ValueError(
              f"Could not parse the precision of `{precision}` to a valid torch.dtype"
          )
  
imports: import os
  import time
  from typing import Callable, Dict, List, Optional, Union
  import numpy
  import torch
  import torch.nn.functional
  from matplotlib import pyplot
  from sklearn.metrics import classification_report, confusion_matrix
  from torch import Tensor
  from nemo.collections.nlp.modules.common.megatron.utils import erf_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import openai_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import squared_relu
  from nemo.utils import logging
label_desc: Use image node_id 1 for calling a global function with example usage: torch_dtype_from_precision(precision, megatron_amp_O2) and returns: torch, torch, torch, torch