code: def train_for_checkpoint(
      hparams: typing.Dict,
      trial_class: pytorch.PyTorchTrial,
      tmp_path: pathlib.Path,
      exp_config: typing.Dict,
      slots_per_trial: int = 1,
      steps: int = 1,
      trial_args: typing.Optional[typing.Dict] = None,
  ) -> int:
      checkpoint_dir = str(tmp_path.joinpath("checkpoint"))
      tensorboard_path = tmp_path.joinpath("tensorboard")
  
      trial, trial_controller = create_trial_and_trial_controller(
          trial_class=trial_class,
          hparams=hparams,
          slots_per_trial=slots_per_trial,
          exp_config=exp_config,
          max_batches=steps,
          min_validation_batches=steps,
          min_checkpoint_batches=steps,
          checkpoint_dir=checkpoint_dir,
          tensorboard_path=tensorboard_path,
          expose_gpus=True,
          trial_args=trial_args,
      )
  
      trial_controller.run()
  
      assert (
          len(os.listdir(checkpoint_dir)) == 1
      ), "trial did not create a checkpoint"
  
      return trial_controller.state.batches_trained
  
file_path: determined/harness/tests/experiment/pytorch_utils.py
imports: import os
  import pathlib
  import random
  import sys
  import typing
  import uuid
  from torch.distributed import launcher
  import determined
  from determined import gpu, pytorch
  from tests.experiment import utils
label_desc: Use image node_id 2 for calling a global function with example usage: train_for_checkpoint(hparams, trial_class, tmp_path, exp_config, slots_per_trial, steps, trial_args) and returns: trial_controller