code: def test_disable_autocast_linear(tmpdir, half_op):
      amp = pytest.importorskip("torch.cuda.amp")
  
      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()
  
      with amp.autocast(False):
          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 2 for calling a global function with example usage: test_disable_autocast_linear(tmpdir, half_op) without return types