@Sara_Rojas_Martinez I ended up with this code that comes with helper debug functions. Not sure if its correct as my experiments ended up learning two distribution equally well and not have a compounded advantage as shown in the paper.
class MyBatchNorm2d(nn.Module):
    def __init__(self, init_bn, isAuxBN):
        super(MyBatchNorm2d, self).__init__()
        self.bn = copy.deepcopy(init_bn)
        
        # Aux BN
        self.bn_aux = copy.deepcopy(init_bn)
        
        self.isAuxBN = isAuxBN
        self.isForwardAuxBN = False
        
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        
    def forward(self, input):            
        if self.training:
            self.num_batches_tracked = self.num_batches_tracked + 1            
        
        # Odd runs are AUX input
        if self.isAuxBN:
            if self.training and self.num_batches_tracked % 2 != 0:
                return self.bn_aux(input)
            
            if self.isForwardAuxBN:
                return self.bn_aux(input)
        
        return self.bn(input)
    
    def run_stats(self):
        return '\n\nParent {} \t bn {} \t bn_aux {}\n\n'.format(self.num_batches_tracked, self.bn.num_batches_tracked, self.bn_aux.num_batches_tracked)
        
    def bn_stats(self):
        return 'BN mean\n {}, \n\n'.format(self.bn.running_mean)
    
    def bn_aux_stats(self):
        return 'BN-AUX mean\n {}, \n'.format(self.bn_aux.running_mean)
    
    def __repr__(self):
        parent_string = super().__repr__()
        txt = parent_string + self.run_stats()
        txt = txt + self.bn_stats()
        txt = txt + self.bn_aux_stats() 
        return txt + '(is forward using AUX {})'.format(self.isForwardAuxBN)