code: def run_inference(
      ml_model: nn.Module,
      data_loader: Any,
      context: core.Context,
      rank: int,
      skip: int,
      max_shard_length: int,
      pred_dir: pathlib.Path,
      checkpoint_interval: int,
  ) -> None:
      ml_model.eval()
      dataloader_iterator = iter(data_loader)
  
      with torch.no_grad():
          last_checkpoint_step = skip
          steps_completed = skip
          for batch_idx in range(skip, max_shard_length):
              data, label = next(dataloader_iterator, (None, None))
              if data is not None:
                  output = ml_model(data)
                  preds = output.argmax(dim=1, keepdim=True)
  
                  file_name = f"inference_out_{rank}_{batch_idx}.json"
  
                  output = []
  
                  for pred in preds:
                      output.append(pred[0].item())
  
                  with open(
                      os.path.join(pred_dir, f"{file_name}"), "w"
                  ) as f:
                      json.dump({"predictions": output}, f)
  
                  steps_completed = batch_idx + 1
                  logging.info(f"Completed step {steps_completed}")
  
              if steps_completed % checkpoint_interval == 0:
                  checkpoint(steps_completed, context)
                  last_checkpoint_step = steps_completed
                  if context.preempt.should_preempt():
                      return
  
          if steps_completed > last_checkpoint_step:
              checkpoint(steps_completed, context)
  
imports: import json
  import logging
  import math
  import os
  import pathlib
  from typing import Any, Optional
  import constants
  import filelock
  import model
  import torch
  import torch.distributed
  import torchvision
  from torch import nn
  from torch.utils import data
  from torchvision import transforms
  import determined
  from determined import core, pytorch
label_desc: Use image node_id 1 for calling a global function with example usage: run_inference(ml_model, data_loader, context, rank, skip, max_shard_length, pred_dir, checkpoint_interval) without return types