Conditional Batch Normalization was proposed recently and a few recent work seems to suggest this has some interesting properties and give good performance in certain tasks.  In this 341 work, the authors implemented a variant of conditional BN in Tensorflow which learns a different scale and shift for each class.
I tried to implement it in PyTorch with the following code:
import torch
from torch.nn.modules.module import Module
from torch.nn.parameter import Parameter
from torch.nn import functional as F


# TODO: check contiguous in THNN
# TODO: use separate backend functions?
class _CondBatchNorm(Module):

    def __init__(self, num_features, n_labels, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super(_CondBatchNorm, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats
        if self.affine:
            self.weight = Parameter(torch.Tensor(n_labels, num_features))
            self.bias = Parameter(torch.Tensor(n_labels, num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        if self.track_running_stats:
            self.register_buffer('running_mean', torch.zeros(num_features))
            self.register_buffer('running_var', torch.ones(num_features))
        else:
            self.register_parameter('running_mean', None)
            self.register_parameter('running_var', None)
        self.reset_parameters()
        print(self.weight.size())

    def reset_parameters(self):
        if self.track_running_stats:
            self.running_mean.zero_()
            self.running_var.fill_(1)
        if self.affine:
            self.weight.data.uniform_()
            self.bias.data.zero_()

    def _check_input_dim(self, input):
        return NotImplemented

    def forward(self, input, labels):
        self._check_input_dim(input)
        self.weight_per_sample= F.embedding(labels.long(), self.weight)
        self.bias_per_sample = F.embedding(labels.long(), self.bias)
        return F.batch_norm(
            input, self.running_mean, self.running_var, self.weight_per_sample, self.bias_per_sample ,
            self.training or not self.track_running_stats, self.momentum, self.eps)

    def __repr__(self):
        return ('{name}({num_features}, eps={eps}, momentum={momentum},'
                ' affine={affine}, track_running_stats={track_running_stats})'
                .format(name=self.__class__.__name__, **self.__dict__))


However, it seems that the low level torch._C.batch_norm can’t handle a weight/bias matrix (vs. a vector). In comparison, Tensorflow seems to support it well (here 41). I wonder if there is something wrong with the way I implement conditional BN, or Torch/Pytorch currently can’t support it?