code: def test_autocast_linear(self, tmpdir, half_input, half_weight):
      amp = get_accelerator().amp()
  
      hidden_dim = 4
      input = torch.randn(hidden_dim).to(
          get_accelerator().device_name()
      )
      ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).to(
          get_accelerator().device_name()
      )
  
      if half_input:
          input = input.half()
  
      if half_weight:
          ds_linear = ds_linear.half()
  
      with amp.autocast():
          output = ds_linear(input)
          assert (
              output.dtype == torch.half
              or output.dtype == torch.bfloat16
          )
  
file_path: DeepSpeed/tests/unit/runtime/test_autocast.py
imports: import pytest
  import torch
  from deepspeed.runtime.zero.linear import LinearModuleForZeroStage3
  from deepspeed.accelerator import get_accelerator
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling the TestAutoCastEnable obj's underlying member method code with example usage: obj.test_autocast_linear(tmpdir, half_input, half_weight) without return types