code: def _test_mlp_impl(
      self, use_activation: str, bias: bool, enable_autocast: bool
  ):
      mlp = MLP(mlp_sizes, bias=bias, activation=use_activation).cuda()
  
      mlp_layers = []
      for i in range(mlp.num_layers):
          linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=bias)
          with torch.no_grad():
              mlp.weights[i].copy_(linear.weight)
              if bias:
                  mlp.biases[i].copy_(linear.bias)
          mlp_layers.append(linear)
          if use_activation == "relu":
              mlp_layers.append(nn.ReLU())
          if use_activation == "sigmoid":
              mlp_layers.append(nn.Sigmoid())
  
      ref_mlp = nn.Sequential(*mlp_layers).cuda()
  
      test_input = (
          torch.empty(batch_size, mlp_sizes[0], device="cuda")
          .uniform_(-1.0, 1.0)
          .requires_grad_()
      )
      ref_input = test_input.clone().detach().requires_grad_()
  
      with torch.cuda.amp.autocast_mode.autocast(
          enabled=enable_autocast
      ):
          mlp_out = mlp(test_input)
          mlp_loss = mlp_out.mean().mul(10.0)
          # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out
          ref_out = ref_mlp(ref_input)
          ref_loss = ref_out.mean().mul(10.0)
  
      mlp_loss.backward()
      ref_loss.backward()
      if enable_autocast:
          self.assertEqual(mlp_out.dtype, torch.float16)
          self.assertEqual(ref_out.dtype, torch.float16)
      else:
          self.assertEqual(mlp_out, ref_out)
          self.assertEqual(test_input.grad, ref_input.grad)
          self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad)
  
file_path: apex/tests/L0/run_mlp/test_mlp.py
imports: from itertools import product
  from time import time
  import torch
  from torch import nn
  from torch.testing._internal import common_utils
  from torch.testing._internal.common_device_type import instantiate_device_type_tests
  from torch.testing._internal.common_device_type import onlyCUDA
  from apex.mlp import MLP
label_desc: Use image node_id 3 for calling the TestMLP obj's underlying member method code with example usage: obj._test_mlp_impl(use_activation, bias, enable_autocast) without return types