import torch


def normalize_batch_per_example(tensor,
                                normalize_mean=True,
                                normalize_std=True):
  """Normalize batch of tensors by mean and variance

  Mean and variance are computed per channel and example

  Parameters
  ----------
  tensor : torch.Tensor or torch.autograd.Variable
    Tensor of shape BxCxHxW
  norm_mean : bool
    If true, subtract per channel mean
  norm_std : bool
    If true, divide by per channel standard deviation
  """
  b, c, h, w = tensor.shape
  out = tensor.view(b, c, h * w)  # Reshape so that we only need one mean op
  if normalize_mean:
    out = out - torch.mean(out, dim=1, keepdim=True)
  if normalize_std:
    out = out / torch.std(out, dim=1, keepdim=True)
  return out.view(b, c, h, w)