code: def test_correctness(use_fp16, checkpoint, chunks):
      torch.manual_seed(0)
      np.random.seed(0)
  
      if use_fp16 and not hasattr(torch.cuda.amp, "custom_fwd"):
          pytest.skip(
              f"AMP APIs are not supported in torch version {torch.__version__}"
          )
  
      model = _get_model()
      rmodel, ropt, rloss = _train_reg_model(model)
      pmodel, popt, ploss = _train_pipe_model(
          model,
          use_fp16=use_fp16,
          checkpoint=checkpoint,
          chunks=chunks,
      )
      _check_parity(rmodel, pmodel, ropt, popt, rloss, ploss)
  
file_path: fairscale/tests/nn/pipe/test_parity.py
imports: import contextlib
  import copy
  import numpy
  import pytest
  import torch
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu
  from fairscale.nn import Pipe
label_desc: Use image node_id 7 for calling a global function with example usage: test_correctness(use_fp16, checkpoint, chunks) without return types