def METHOD_NAME(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)