code: def _train_single_process_sync_batchnorm(batch_size, num_steps):
      seed_everything(3)
      dataset = torch.arange(64, dtype=torch.float).view(-1, 1)
      train_dataloader = DataLoader(dataset, batch_size=batch_size)
      model = SyncBNModule(batch_size=batch_size)
      optimizer = model.configure_optimizers()
      model.train()
      for batch_idx, batch in enumerate(train_dataloader):
          optimizer.zero_grad()
          loss = model.training_step(batch, batch)
          loss.backward()
          optimizer.step()
          if batch_idx == num_steps - 1:
              break
  
      return torch.stack(model.bn_outputs)
  
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 2 for calling a global function with example usage: _train_single_process_sync_batchnorm(batch_size, num_steps) and returns: torch