fix bn stat

model.py

@@ -38,12 +38,13 @@ class BatchNormVAE(nn.Module): # https://spaces.ac.cn/archives/7381/
         super(BatchNormVAE, self).__init__()
         kwargs['affine'] = False
         self.TAU = 0.5
-        self.bn = nn.BatchNorm1d(num_features, **kwargs)
+        self.bn_mu = nn.BatchNorm1d(num_features, **kwargs)
+        self.bn_sigma = nn.BatchNorm1d(num_features, **kwargs)
         self.theta = nn.Parameter(torch.zeros(1))
 
     def forward(self, mu, sigma):
-        mu = self.bn(mu)
-        sigma = self.bn(sigma)
+        mu = self.bn_mu(mu)
+        sigma = self.bn_sigma(sigma)
         scale_mu = torch.sqrt(self.TAU + (1 - self.TAU) * F.sigmoid(self.theta))
         scale_sigma = torch.sqrt((1 - self.TAU) * F.sigmoid(-self.theta))
         return mu*scale_mu, sigma*scale_sigma