code: import torch.nn as nn
  import torch.nn.functional as F
  from torch.nn import init
  
  
  class ConditionalBatchNorm2d(nn.BatchNorm2d):
  
      """Conditional Batch Normalization"""
  
      def __init__(self, num_features, eps=1e-05, momentum=0.1,
                   affine=False, track_running_stats=True):
          super(ConditionalBatchNorm2d, self).__init__(
              num_features, eps, momentum, affine, track_running_stats
          )
  
      def forward(self, input, weight, bias, **kwargs):
          self._check_input_dim(input)
  
          exponential_average_factor = 0.0
  
          if self.training and self.track_running_stats:
              self.num_batches_tracked += 1
              if self.momentum is None:  # use cumulative moving average
                  exponential_average_factor = 1.0 /\
                       self.num_batches_tracked.item()
              else:  # use exponential moving average
                  exponential_average_factor = self.momentum
  
          output = F.batch_norm(input, self.running_mean, self.running_var,
                                self.weight, self.bias,
                                self.training or not self.track_running_stats,
                                exponential_average_factor, self.eps)
          if weight.dim() == 1:
              weight = weight.unsqueeze(0)
          if bias.dim() == 1:
              bias = bias.unsqueeze(0)
          size = output.size()
          weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)
          bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)
          return weight * output + bias
  
  
  class CategoricalConditionalBatchNorm2d(ConditionalBatchNorm2d):
  
      def __init__(self, num_classes, num_features, eps=1e-5, momentum=0.1,
                   affine=False, track_running_stats=True):
          super(CategoricalConditionalBatchNorm2d, self).__init__(
              num_features, eps, momentum, affine, track_running_stats
          )
          self.weights = nn.Embedding(num_classes, num_features)
          self.biases = nn.Embedding(num_classes, num_features)
  
          self._initialize()
  
      def _initialize(self):
          init.ones_(self.weights.weight.data)
          init.zeros_(self.biases.weight.data)
  
      def forward(self, input, c, **kwargs):
          weight = self.weights(c)
          bias = self.biases(c)
  
          return super(CategoricalConditionalBatchNorm2d, self)\
              .forward(input, weight, bias)
  
  
  if __name__ == '__main__':
      """Forward computation check."""
      import torch
      size = (3, 3, 12, 12)
      batch_size, num_features = size[:2]
      print('# Affirm embedding output')
      naive_bn = nn.BatchNorm2d(3)
      idx_input = torch.tensor([1, 2, 0], dtype=torch.long)
      embedding = nn.Embedding(3, 3)
      weights = embedding(idx_input)
      print('# weights size', weights.size())
      empty = torch.tensor((), dtype=torch.float)
      running_mean = empty.new_zeros((3,))
      running_var = empty.new_ones((3,))
  
      naive_bn_W = naive_bn.weight
  
      input = torch.rand(*size, dtype=torch.float32)
      print('input size', input.size())
      print('input ndim ', input.dim())
  
      _ = naive_bn(input)
  
      print('# batch_norm with given weights')
  
      try:
          with torch.no_grad():
              output = F.batch_norm(input, running_mean, running_var,
                                    weights, naive_bn.bias, False, 0.0, 1e-05)
      except Exception as e:
          print("\tFailed to use given weights")
          print('# Error msg:', e)
          print()
      else:
          print("Succeeded to use given weights")
  
      print('\n# Batch norm before use given weights')
      with torch.no_grad():
          tmp_out = F.batch_norm(input, running_mean, running_var,
                                 naive_bn_W, naive_bn.bias, False, .0, 1e-05)
      weights_cast = weights.unsqueeze(-1).unsqueeze(-1)
      weights_cast = weights_cast.expand(tmp_out.size())
      try:
          out = weights_cast * tmp_out
      except Exception:
          print("Failed")
      else:
          print("Succeeded!")
          print('\t {}'.format(out.size()))
          print(type(tuple(out.size())))
  
      print('--- condBN and catCondBN ---')
  
      catCondBN = CategoricalConditionalBatchNorm2d(3, 3)
      output = catCondBN(input, idx_input)
  
      assert tuple(output.size()) == size
  
      condBN = ConditionalBatchNorm2d(3)
  
      idx = torch.tensor([1], dtype=torch.long)
      out = catCondBN(input, idx)
  
      print('cat cond BN weights\n', catCondBN.weights.weight.data)
      print('cat cond BN biases\n', catCondBN.biases.weight.data)
  
apis: torch.nn.BatchNorm2d torch.rand torch.no_grad torch.tensor torch.nn.Embedding torch.nn.functional.batch_norm torch.nn.init.ones_ torch.nn.init.zeros_