class BatchNorm3dFix(nn.BatchNorm3d):
    """
    In this way, the stat values can be correct now.
    """
    def __init__(self, *args, **kwargs):
        super(BatchNorm3dFix, self).__init__(*args, **kwargs)
        self.prev_running_mean = self.running_mean.new(self.running_mean.size())
        self.prev_running_var = self.running_var.new(self.running_var.size())

    def forward(self, input, in_backward=False):
        self._check_input_dim(input)
        if in_backward:
            self.running_mean.copy_(self.prev_running_mean)
            self.running_var.copy_(self.prev_running_var)
        else:
            self.prev_running_mean.copy_(self.running_mean)
            self.prev_running_var.copy_(self.running_var)

        return F.batch_norm(
            input, self.running_mean, self.running_var, self.weight, self.bias,
            self.training or not self.track_running_stats,
            self.momentum, self.eps)


class _ConvFunction(object):
    def __init__(self, norm, relu, conv):
        """
        function for checkpoint
        Concat(if needed) -> BatchNorm -> ReLU -> Conv
        """
        self.norm = norm
        self.relu = relu
        self.conv = conv
        self.in_backward = False

    def __call__(self, *inputs):
        concated_features = torch.cat(inputs, 1) if len(inputs) > 1 else inputs[0]
        output = self.conv(self.relu(self.norm(concated_features, self.in_backward)))
        self.in_backward = True
        return output