code: def test_pipe_use_reentrant(self, topo_config):
      skip_on_arch(min_arch=7)
  
      topo = PipeTopo(**topo_config)
      steps = 100  # must be >=100
  
      # Allocate model for consistent initial weights.
      init_net = AlexNetPipe()
  
      # Train with not set use_reentrant, default: True
      base_net = copy.deepcopy(init_net)
      base_model = PipelineModule(
          layers=base_net.to_layers(),
          topology=topo,
          loss_fn=nn.CrossEntropyLoss(),
      )
      base_losses = train_cifar(
          base_model,
          config=config_dict,
          num_steps=steps,
          fp16=config_dict["fp16"]["enabled"],
      )
  
      # Train with set use_reentrant=False, this will use ``non_reentrant_checkpoint``
      test_config_dict = copy.deepcopy(config_dict)
      test_config_dict["pipeline"]["use_reentrant"] = False
      test_net = copy.deepcopy(init_net)
      test_model = PipelineModule(
          layers=test_net.to_layers(),
          topology=topo,
          loss_fn=nn.CrossEntropyLoss(),
      )
      test_losses = train_cifar(
          test_model,
          config=test_config_dict,
          num_steps=steps,
          fp16=config_dict["fp16"]["enabled"],
      )
  
      abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)]
      rel_diffs = [
          rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)
      ]
      if dist.get_rank() == 0:
          print(
              f"abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}"
          )
          print(
              f"rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}"
          )
          print(
              f"first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}"
          )
  
          for lastX in [1, 10, 100]:
              base_avg = sum(base_losses[-lastX:]) / lastX
              test_avg = sum(test_losses[-lastX:]) / lastX
              print(
                  f"last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}"
              )
      lastX = 100
      base = base_losses[-lastX:]
      base_avg = sum(base) / len(base)
      test = test_losses[-lastX:]
      test_avg = sum(test) / len(test)
      assert rel_diff(base_avg, test_avg) < 0.05
  
file_path: deepspeed/tests/unit/runtime/pipe/test_pipe.py
imports: import copy
  import torch.nn
  import pytest
  import deepspeed.comm
  from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
  from deepspeed.runtime.pipe.module import PipelineModule
  from unit.alexnet_model import AlexNetPipe, train_cifar
  from unit.common import DistributedTest
  from unit.util import skip_on_arch
label_desc: Use image node_id 2 for calling the TestPipeCifar10 obj's underlying member method code with example usage: obj.test_pipe_use_reentrant(topo_config) without return types