code: def test_single_run():
      device, offload_device = _init()
      model = _get_model()
  
      peak_mem = {}
      for checkpoint_activation in [True, False]:
          offload_model = OffloadModel(
              model=model,
              device=device,
              offload_device=offload_device,
              num_slices=2,
              checkpoint_activation=checkpoint_activation,
          )
          offload_optimizer = torch.optim.SGD(
              offload_model.parameters(), lr=0.001
          )
  
          input = torch.ones(1000, 2).to(device)
          labels = torch.ones(1000, 2).to(device)
          offload_model.train()
          pred = offload_model(input)
          loss_fn = torch.nn.MSELoss(reduction="sum")
          loss = loss_fn(pred, labels)
          loss.backward()
          offload_optimizer.step()
          key = "ca_" + str(checkpoint_activation)
          peak_mem[key] = torch.cuda.memory_stats(0)[
              "allocated_bytes.all.peak"
          ]
          print(
              "Peak allocated bytes on cuda:0 for checkpoint_activation "
              + str(checkpoint_activation)
              + ": {:2f}".format(peak_mem[key])
          )
  
      # TODO(anj-s): We need a better requirement since this fails on CircleCI right now.
      assert peak_mem["ca_True"] <= peak_mem["ca_False"]
  
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 2 for calling a global function with example usage: test_single_run() without return types