Merge pull request #15 from y0ast/patch-1

Properly register parameters of LU decomposition

Author: chaiyujin

glow/modules.py

@@ -208,8 +208,8 @@ def __init__(self, num_channels, LU_decomposed=False):
             l_mask = np.tril(np.ones(w_shape, dtype=np.float32), -1)
             eye = np.eye(*w_shape, dtype=np.float32)
 
-            self.p = torch.Tensor(np_p.astype(np.float32))
-            self.sign_s = torch.Tensor(np_sign_s.astype(np.float32))
+            self.register_buffer('p', torch.Tensor(np_p.astype(np.float32)))
+            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(np.float32)))
             self.l = nn.Parameter(torch.Tensor(np_l.astype(np.float32)))
             self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(np.float32)))
             self.u = nn.Parameter(torch.Tensor(np_u.astype(np.float32)))