code: def checkpoint_eval(pipe_class):
      model = nn.Sequential(nn.Linear(1, 1))
      model = pipe_class(
          model,
          balance=[1],
          worker_map=get_worker_map(),
          chunks=2,
      )
      input = torch.rand(2, 1)
  
      def find_grad_fn(grad_fn, name):
          if grad_fn is None:
              return False
          if grad_fn.__class__.__name__ == name:
              return True
          for next_grad_fn, _ in grad_fn.next_functions:
              if find_grad_fn(next_grad_fn, name):
                  return True
          return False
  
      model.train()
      train_output = model(input)
      assert find_grad_fn(train_output.grad_fn, "CheckpointBackward")
      assert find_grad_fn(train_output.grad_fn, "RecomputeBackward")
  
      model.eval()
      eval_output = model(input)
      assert not find_grad_fn(eval_output.grad_fn, "CheckpointBackward")
      assert not find_grad_fn(eval_output.grad_fn, "RecomputeBackward")
  
imports: from collections import OrderedDict
  from copy import deepcopy
  import os
  import time
  import pytest
  import torch
  from torch import nn
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.internal import torch_version
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from fairscale.nn.pipe import AsyncPipe
  from fairscale.nn.pipe.types import LazyModule
label_desc: Use image node_id 17 for calling a global function with example usage: checkpoint_eval(pipe_class) and returns: False, False, True, True