code: def run_amp(
      tmp_path: pathlib.Path,
      api_style: str,
      batches_trained: typing.Optional[int] = 0,
  ):
      checkpoint_dir = str(tmp_path.joinpath("checkpoint"))
      class_selector = {
          "apex": "MNistApexAMPTrial",
          "auto": "MNistAutoAMPTrial",
          "manual": "MNistManualAMPTrial",
      }
  
      config = utils.load_config(
          utils.fixtures_path(
              f"pytorch_amp/{api_style}_amp_distributed.yaml"
          )
      )
      config = config.copy()
      config.setdefault("profiling", {})
      config["profiling"]["enabled"] = True
  
      hparams = config["hyperparameters"]
  
      exp_config = utils.make_default_exp_config(
          hparams,
          scheduling_unit=1,
          searcher_metric="validation_loss",
          checkpoint_dir=checkpoint_dir,
      )
      exp_config.update(config)
      exp_config["searcher"]["smaller_is_better"] = True
  
      example_path = utils.fixtures_path(
          f"pytorch_amp/{api_style}_amp_model_def.py"
      )
      trial_class = utils.import_class_from_module(
          class_selector[api_style], example_path
      )
      trial_class._searcher_metric = "validation_loss"
  
      if batches_trained == 0:
          return pytorch_utils.train_for_checkpoint(
              trial_class=trial_class,
              hparams=hparams,
              slots_per_trial=2,
              tmp_path=tmp_path,
              exp_config=exp_config,
              steps=1,
              trial_args={"hparams": hparams},
          )
      else:
          pytorch_utils.train_from_checkpoint(
              trial_class=trial_class,
              hparams=hparams,
              slots_per_trial=2,
              tmp_path=tmp_path,
              exp_config=exp_config,
              steps=(1, 1),
              batches_trained=batches_trained,
              trial_args={"hparams": hparams},
          )
          return True
  
imports: import contextlib
  import importlib
  import io
  import os
  import pathlib
  import sys
  import typing
  from unittest import mock
  import numpy
  import pytest
  import torch
  from _pytest import monkeypatch
  from torch.distributed import launcher
  import determined
  from determined import pytorch
  from tests.experiment import pytorch_utils, utils
  from tests.experiment.fixtures import pytorch_onevar_model
label_desc: Use image node_id 5 for calling a global function with example usage: run_amp(tmp_path, api_style, batches_trained) and returns: pytorch_utils, True