fix bugs

ELIFE-ASU · Apr 29, 2021 · 9c7e2f4 · 9c7e2f4
1 parent 84a1277
commit 9c7e2f4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/INN/INN.py b/src/INN/INN.py
@@ -284,7 +284,7 @@ def forward(self, x, log_p0=0, log_det_J=0):
             log_det = self.logdet(x)
         if len(x.shape) == 2:
             # [batch, dim]
-            log_det = self.logdet(x).repeat(x.shape[0])
+            log_det = self.logdet(x)
         x = super(Linear, self).forward(x)
 
         if self.compute_p:

diff --git a/src/INN/utilities.py b/src/INN/utilities.py
@@ -479,8 +479,9 @@ def init_weights(self, m):
             nonlinearity = 'sigmoid'
 
         if type(m) == nn.Linear:
-            # doing Kaiming initialization
-            torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity=nonlinearity)
+            # doing xavier initialization
+            # NOTE: Kaiming initialization will make the output too high, which leads to nan
+            torch.nn.init.xavier_uniform_(m.weight.data)
             torch.nn.init.zeros_(m.bias.data)
 
     def forward(self, x):