code: def test_correctness(
      use_fp16, checkpoint_activation, num_microbatches, use_auto_shard
  ):
      pytest.skip("skip this test until the issue #900 is resolved.")
      if use_auto_shard and torch_version() < (1, 8, 0):
          pytest.skip("auto_shard requires torch version >= 1.8.0")
  
      if (use_fp16 or checkpoint_activation) and not hasattr(
          torch.cuda.amp, "custom_fwd"
      ):
          pytest.skip(
              f"AMP APIs are not supported in torch version {torch.__version__}"
          )
  
      if not checkpoint_activation and num_microbatches > 1:
          pytest.skip(
              "We only support microbatches with activation offloading."
          )
  
      device, offload_device = _init()
      model = _get_model()
      if use_auto_shard:
          offload_model = shard_model(model)
      else:
          offload_model = model
  
      rmodel, ropt, rloss = _train_reg_model(
          model, device, offload_device
      )
      omodel, oopt, oloss = _train_offload_model(
          offload_model,
          device,
          offload_device,
          use_fp16=use_fp16,
          checkpoint_activation=checkpoint_activation,
          num_microbatches=num_microbatches,
      )
      _check_parity(
          rmodel.cpu(), omodel.cpu(), ropt, oopt, rloss, oloss
      )
  
file_path: fairscale/tests/experimental/nn/test_offload.py
imports: import contextlib
  import copy
  import numpy
  import pytest
  import torch
  from fairscale.experimental.nn.offload import OffloadModel
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
  from fairscale.internal import torch_version
label_desc: Use image node_id 9 for calling a global function with example usage: test_correctness(use_fp16, checkpoint_activation, num_microbatches, use_auto_shard) without return types