code: def test_missing_amp_autocast(tmpdir, half_op):
      hidden_dim = 4
      if half_op:
          input = torch.randn(hidden_dim).cuda().half()
          ds_linear = (
              LinearModuleForZeroStage3(hidden_dim, hidden_dim)
              .cuda()
              .half()
          )
      else:
          input = torch.randn(hidden_dim).cuda()
          ds_linear = LinearModuleForZeroStage3(
              hidden_dim, hidden_dim
          ).cuda()
  
      output = ds_linear(input)
      assert output.dtype == ds_linear.weight.dtype
  
imports: import pytest
  import torch
  from deepspeed.runtime.zero.linear import LinearModuleForZeroStage3
label_desc: Use image node_id 1 for calling a global function with example usage: test_missing_amp_autocast(tmpdir, half_op) without return types