title: How to implement Batchnorm2d in Pytorch myself?
body: <p>I'm trying to implement Batchnorm2d() layer with:</p>
  <pre><code>class BatchNorm2d(nn.Module):
  
      def __init__(self, num_features):
          super(BatchNorm2d, self).__init__()
          self.num_features = num_features
          device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
          self.eps = 1e-5
          self.momentum = 0.1
          self.first_run = True
  
      def forward(self, input):
          # input: [batch_size, num_feature_map, height, width]
          device = input.device
          if self.training:
              mean = torch.mean(input, dim=0, keepdim=True).to(device)  # [1, num_feature, height, width]
              var = torch.var(input, dim=0, unbiased=False, keepdim=True).to(device)  # [1, num_feature, height, width]
              if self.first_run:
                  self.weight = Parameter(torch.randn(input.shape, dtype=torch.float32, device=device), requires_grad=True)
                  self.bias = Parameter(torch.randn(input.shape, dtype=torch.float32, device=device), requires_grad=True)
                  self.register_buffer('running_mean', torch.zeros(input.shape).to(input.device))
                  self.register_buffer('running_var', torch.ones(input.shape).to(input.device))
                  self.first_run = False
              self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
              self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
              bn_init = (input - mean) / torch.sqrt(var + self.eps)
          else:
              bn_init = (input - self.running_mean) / torch.sqrt(self.running_var + self.eps)
          return self.weight * bn_init + self.bias
  </code></pre>
  <p>But after training &amp; testing I found that the results using my layer is incomparable with the results using <code>nn.Batchnorm2d()</code>. There must be something wrong with it, and I guess the problem relates to initializing parameters in <code>forward()</code>? I did that because I don't know how to know the shape of input in <code>__init__()</code>, maybe there is a better way.  I don't know how to fix it, please help. Thanks!!</p>