code: def test_sync_batchnorm_parity(tmpdir):
      """Test parity between 1) Training a synced batch-norm layer on 2 GPUs with batch size B per device 2) Training a
      batch-norm layer on CPU with twice the batch size."""
      seed_everything(3)
      # 2 GPUS, batch size = 4 per GPU => total batch size = 8
      model = SyncBNModule(batch_size=4)
      trainer = Trainer(
          default_root_dir=tmpdir,
          accelerator="gpu",
          strategy="ddp_find_unused_parameters_true",
          devices=2,
          max_steps=3,
          sync_batchnorm=True,
          num_sanity_val_steps=0,
          use_distributed_sampler=False,
          deterministic=True,
          benchmark=False,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
  
      # the strategy is responsible for tearing down the batch norm wrappers
      assert not isinstance(
          model.bn_layer, torch.nn.modules.batchnorm.SyncBatchNorm
      )
      assert isinstance(
          model.bn_layer, torch.nn.modules.batchnorm._BatchNorm
      )
  
      bn_outputs = torch.stack(
          model.bn_outputs
      )  # 2 x 4 x 1 on each GPU
      bn_outputs_multi_device = trainer.strategy.all_gather(
          bn_outputs
      ).cpu()  # 2 x 2 x 4 x 1
  
      if trainer.global_rank == 0:
          # pretend we are now training on a single GPU/process
          # (we are reusing the rank 0 from the previous training)
  
          # 1 GPU, batch size = 8 => total batch size = 8
          bn_outputs_single_device = (
              _train_single_process_sync_batchnorm(
                  batch_size=8, num_steps=3
              )
          )
  
          gpu0_outputs = bn_outputs_multi_device[0]  # 2 x 4 x 1
          gpu1_outputs = bn_outputs_multi_device[1]  # 2 x 4 x 1
          slice0 = bn_outputs_single_device[:, 0::2]
          slice1 = bn_outputs_single_device[:, 1::2]
  
          assert torch.allclose(gpu0_outputs, slice0)
          assert torch.allclose(gpu1_outputs, slice1)
  
file_path: pytorch-lightning/tests/parity_pytorch/test_sync_batchnorm_parity.py
imports: import torch
  import torch.nn
  from lightning.pytorch import LightningModule, Trainer, seed_everything
  from torch.utils.data import DataLoader, DistributedSampler
  from parity_pytorch import RunIf
label_desc: Use image node_id 1 for calling a global function with example usage: test_sync_batchnorm_parity(tmpdir) without return types