code: def test_offload_memory():
      if torch_version() >= (1, 12, 0):
          pytest.skip("to be fixed")
  
      device = "cuda"
  
      input = torch.rand(60, 24, 4).requires_grad_(True)
  
      model = CpuOffloadModel().to(device)
      base = get_loss_and_gnorm(model, input.to(device))
  
      model = CpuOffloadModel(True).to(device)
      cpt = get_loss_and_gnorm(model, input.to(device))
  
      model = CpuOffloadModel(True, True).to(device)
      offload = get_loss_and_gnorm(model, input.to(device))
  
      for key in "loss", "gnorm":
          if not (base[key] == cpt[key] == offload[key]):
              # Use print to collect all debugging info.
              print(base, cpt, offload)
              assert 0
          del base[key]
          del cpt[key]
          del offload[key]
  
      ref_base = {
          "mem_0": 32256,
          "mem_peak": 334336,
          "mem_after_fwd": 274944,
          "mem_after_bwd": 41984,
      }
      ref_cpt = {
          "mem_0": 32256,
          "mem_peak": 253952,
          "mem_after_fwd": 101888,
          "mem_after_bwd": 41984,
      }
      ref_offload = {
          "mem_0": 32256,
          "mem_peak": 207872,
          "mem_after_fwd": 55808,
          "mem_after_bwd": 41984,
      }
  
      if not (
          base == ref_base and cpt == ref_cpt and offload == ref_offload
      ):
          # Use print to collect all debugging info.
          print(base, cpt, offload)
          assert 0
  
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 4 for calling a global function with example usage: test_offload_memory() without return types