code: def test_autocast_linear(tmpdir, half_input, half_weight):
      amp = pytest.importorskip("torch.cuda.amp")
  
      hidden_dim = 4
      input = torch.randn(hidden_dim).cuda()
      ds_linear = LinearModuleForZeroStage3(
          hidden_dim, hidden_dim
      ).cuda()
  
      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
  
imports: import pytest
  import torch
  from deepspeed.runtime.zero.linear import LinearModuleForZeroStage3
label_desc: Use image node_id 3 for calling a global function with example usage: test_autocast_linear(tmpdir, half_input, half_weight) without return types