class_docstr: Fast re-implementation of ``torch.nn.SyncBatchNorm`` that can achieve a speedup
  of 5x or more over the default implementation depending on size of the input
  and number of distributed workers.
class_docstr_tok: ["Fast","re-implementation","of","``","torch.nn.SyncBatchNorm","``","that","can","achieve","a","speedup","of","5x","or","more","over","the","default","implementation","depending","on","size","of","the","input","and","number","of","distributed","workers","."]
code: def forward(self, input: Tensor) -> Tensor:  # type: ignore
      # There are 3 modes this is being called:
      # 1. not wrapped (and there is only a single phase)
      # 2. wrapped and in checkpointing phase
      # 3. wrapped and in recomputing phase
  
      if not dist.is_initialized() or not self.training:
          return super().forward(input)
  
      wrapped = is_checkpointing() or is_recomputing()
      if not wrapped or is_checkpointing():
          # NOTE The full backward, including mean and var, is done by _SyncBatchNormFunction.
          with torch.no_grad():
              mean, var, invstd, total_count = _calculate_stats(
                  input, self.eps, self._process_group
              )
              if self.track_running_stats:
                  _track_running_stats(
                      self.running_mean,
                      self.running_var,
                      self.momentum,
                      mean,
                      var,
                      total_count,
                  )
  
      if is_checkpointing():
          self.saved_for_2nd_fwd.append((mean, invstd, total_count))
          return _forward(
              input, self.affine, mean, invstd, self.weight, self.bias
          )
      if is_recomputing():
          mean, invstd, total_count = self.saved_for_2nd_fwd.pop(0)
  
      return _SyncBatchNormFunction.apply(
          input,
          self.weight,
          self.bias,
          self.affine,
          mean,
          invstd,
          total_count,
          self._process_group,
      )
  
file_path: fairscale/fairscale/experimental/nn/sync_batchnorm.py
imports: from typing import Any, Dict, List, Optional, Tuple
  import torch
  from torch import Tensor
  import torch.distributed
  from torch.distributed import ProcessGroup
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint import is_checkpointing, is_recomputing
label_desc: Use image node_id 2 for calling the SyncBatchNorm obj's underlying member method code with example usage: obj.forward(input) and returns: _SyncBatchNormFunction, super, _forward