code: def test(
      self,
      model_w_task,
      query,
      inf_kwargs,
      assert_fn,
      dtype,
  ):
      invalid_test_msg = validate_test(
          model_w_task,
          dtype,
          enable_cuda_graph=False,
          enable_triton=False,
      )
      if invalid_test_msg:
          pytest.skip(invalid_test_msg)
  
      if dtype not in get_accelerator().supported_dtypes():
          pytest.skip(
              f"Acceleraor {get_accelerator().device_name()} does not support {dtype}."
          )
  
      # TODO: enable this test after torch 2.1 stable release
      if (
          dtype == torch.bfloat16
          and model_w_task[0] == "Salesforce/codegen-350M-mono"
      ):
          pytest.skip(
              "Codegen model(bf16) need to use torch version > 2.0."
          )
  
      model, task = model_w_task
      local_rank = int(os.getenv("LOCAL_RANK", "0"))
      world_size = int(os.getenv("WORLD_SIZE", "2"))
  
      # We have to load these large models on CPU with pipeline because not
      # enough GPU memory
      pipe = pipeline(
          task, model=model, device=torch.device("cpu"), framework="pt"
      )
      bs_output = pipe(query, **inf_kwargs)
  
      pipe.model = deepspeed.init_inference(
          pipe.model, mp_size=world_size, dtype=dtype
      )
      # Switch device to GPU so that input tensors are not on CPU
      pipe.device = torch.device(
          get_accelerator().device_name(local_rank)
      )
      ds_output = pipe(query, **inf_kwargs)
  
      print(local_rank, "baseline", bs_output)
      print(local_rank, "deepspeed", ds_output)
      assert assert_fn(bs_output, ds_output)
  
imports: import os
  import time
  import torch
  import pytest
  import itertools
  import deepspeed
  from deepspeed.git_version_info import torch_info
  from unit.common import DistributedTest
  from packaging import version
  from deepspeed.ops.op_builder import OpBuilder
  from transformers import pipeline, AutoTokenizer
  from transformers.models.t5.modeling_t5 import T5Block
  from transformers.models.roberta.modeling_roberta import RobertaLayer
  from huggingface_hub import HfApi
  from deepspeed.model_implementations import DeepSpeedTransformerInference
  from torch import nn
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import InferenceBuilder
label_desc: Use image node_id 1 for calling the TestAutoTensorParallelism obj's underlying member method code with example usage: obj.test(model_w_task, query, inf_kwargs, assert_fn, dtype) without return types