86f32d7fc3ee4afe67f6d0bd2a2a8db042406987,pixyz/flows/operations.py,BatchNormFlow,__init__,#BatchNormFlow#Any#Any#,224

Before Change

        self._running_mean = torch.zeros(in_channels, 1, 1)
        self._running_var = torch.ones(in_channels, 1, 1)
        self._batch_mean = torch.zeros(in_channels, 1, 1)
        self._batch_var = torch.ones(in_channels, 1, 1)

    def forward(self, x, compute_jacobian=True):
        if self.training:

After Change

        self.register_buffer("_running_mean", torch.zeros(in_features))
        self.register_buffer("_running_var", torch.ones(in_features))
        self.register_buffer("_batch_mean", torch.zeros(in_features))
        self.register_buffer("_batch_var", torch.ones(in_features))

    def forward(self, x, compute_jacobian=True):
        if self.training:
            self._batch_mean = x.mean(0)

Instances