code: def deferred_batch_norm(checkpoint, lazy, pipe_class):
      bn = nn.BatchNorm2d(3)
      pipe_bn = deepcopy(bn)
      pipe_fn = lambda: pipe_bn  # noqa: E731
      if lazy:
          model = [LazyModule(pipe_fn)]
      else:
          model = nn.Sequential(pipe_bn)
      pipe = pipe_class(
          model,
          balance=[1],
          worker_map=get_worker_map(),
          chunks=2,
          checkpoint=checkpoint,
          deferred_batch_norm=True,
      )
  
      x = torch.rand(4, 3, 10, 10)
      pipe(x).mean().backward()
      bn(x).mean().backward()
  
      assert torch.allclose(
          pipe[0].running_mean, bn.running_mean, atol=1e-4
      )
      assert torch.allclose(
          pipe[0].running_var, bn.running_var, atol=1e-4
      )
  
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 27 for calling a global function with example usage: deferred_batch_norm(checkpoint, lazy, pipe_class) without return types