docstr: Train a linear model using Determined with Automated Mixed Precision in three ways:
  Using Apex and using PyTorch AMP both "automatically" and "manually". In the "manual" case,
  we use the context manager ``autoscale`` in the model's training and
  evaluating methods; a scaler object is wrapped in a Determined context. The same
  is done under the hood in the first two cases.
code: def test_amp(self, tmp_path: pathlib.Path, trial_class) -> None:
      """Train a linear model using Determined with Automated Mixed Precision in three ways:
      Using Apex and using PyTorch AMP both "automatically" and "manually". In the "manual" case,
      we use the context manager ``autoscale`` in the model's training and
      evaluating methods; a scaler object is wrapped in a Determined context. The same
      is done under the hood in the first two cases.
      """
      tensorboard_path = tmp_path.joinpath("tensorboard")
  
      if (
          trial_class is pytorch_onevar_model.OneVarApexAMPTrial
          and not HAVE_APEX
      ):
          pytest.skip("Apex not available")
  
      # The assertions logic in make_amp_workloads require a batch size of one
      hparams = dict(self.hparams)
      hparams["global_batch_size"] = 1
  
      (
          trial,
          controller,
      ) = pytorch_utils.create_trial_and_trial_controller(
          trial_class=trial_class,
          hparams=hparams,
          trial_seed=self.trial_seed,
          expose_gpus=True,
          max_batches=20,
          min_validation_batches=1,
          min_checkpoint_batches=sys.maxsize,
          tensorboard_path=tensorboard_path,
      )
  
      controller.run()
  
      metrics_callback = trial.metrics_callback
      training_metrics = metrics_callback.training_metrics
  
      amp_metrics_test(trial_class, training_metrics)
  
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 21 for calling the TestPyTorchTrial obj's underlying member method code with example usage: obj.test_amp(tmp_path, trial_class) without return types