code: def test_basic(device):
      if "cuda" in device and not torch.cuda.is_available():
          pytest.skip("test requires a GPU")
  
      input = torch.rand(2, 16, 32).requires_grad_(True)
      model = BasicModel().to(device)
      no_cpt = get_loss_and_gnorm(model, input.to(device))
  
      model = BasicModel(use_pytorch_checkpoint=True).to(device)
      pyt_cpt = get_loss_and_gnorm(model, input.to(device))
  
      model = BasicModel(use_fairscale_checkpoint=True).to(device)
      fairscale_cpt = get_loss_and_gnorm(model, input.to(device))
  
      model = BasicModel(
          use_fairscale_checkpoint=True, offload_to_cpu=True
      ).to(device)
      fairscale_cpt_offload = get_loss_and_gnorm(
          model, input.to(device)
      )
  
      # Check for correctness.
      for key in "loss", "gnorm":
          if not (
              no_cpt[key]
              == pyt_cpt[key]
              == fairscale_cpt[key]
              == fairscale_cpt_offload[key]
          ):
              print(
                  no_cpt, pyt_cpt, fairscale_cpt, fairscale_cpt_offload
              )
              assert 0
          del no_cpt[key]
          del pyt_cpt[key]
          del fairscale_cpt[key]
          del fairscale_cpt_offload[key]
  
      # Check for memory usage for cuda only.
      if "cpu" in device:
          return
  
      mem_peaks = [98816, 103424, 103424, 107520]
      if torch_version() < (1, 7, 0):
          # Older torch behaves slightly differently
          mem_peaks = [102400, 103424, 103424, 107520]
  
      assert no_cpt == {
          "mem_0": 38912,
          "mem_peak": mem_peaks[0],
          "mem_after_fwd": 64000,
          "mem_after_bwd": 74240,
      }, no_cpt
      assert pyt_cpt == {
          "mem_0": 38912,
          "mem_peak": mem_peaks[1],
          "mem_after_fwd": 43520,
          "mem_after_bwd": 74240,
      }, pyt_cpt
      assert fairscale_cpt == {
          "mem_0": 38912,
          "mem_peak": mem_peaks[2],
          "mem_after_fwd": 43520,
          "mem_after_bwd": 74240,
      }, fairscale_cpt
      assert fairscale_cpt_offload == {
          "mem_0": 38912,
          "mem_peak": mem_peaks[3],
          "mem_after_fwd": 43520,
          "mem_after_bwd": 74240,
      }, fairscale_cpt_offload
  
imports: import pytest
  import torch
  import torch.nn
  from torch.utils.checkpoint import checkpoint
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper, disable_checkpointing
  from fairscale.nn.misc import FlattenParamsWrapper
  from fairscale.nn.misc import checkpoint_wrapper
label_desc: Use image node_id 3 for calling a global function with example usage: test_basic(device) without return types