Merge pull request #171 from tianjiao-j/master

fixed error caused by in-place operations

FrEIA/modules/coupling_layers.py

@@ -382,7 +382,7 @@ def _coupling1(self, x1, u2, rev=False):
         a2 = self.subnet2(u2)
         s2, t2 = a2[:, :self.split_len1], a2[:, self.split_len1:]
         s2 = self.clamp * self.f_clamp(s2)
-        s2 -= s2.mean(1, keepdim=True)
+        s2 = s2 - s2.mean(1, keepdim=True)  # changed inplace op
 
         if rev:
             y1 = (x1 - t2) * torch.exp(-s2)
@@ -395,7 +395,7 @@ def _coupling2(self, x2, u1, rev=False):
         a1 = self.subnet1(u1)
         s1, t1 = a1[:, :self.split_len2], a1[:, self.split_len2:]
         s1 = self.clamp * self.f_clamp(s1)
-        s1 -= s1.mean(1, keepdim=True)
+        s1 = s1 - s1.mean(1, keepdim=True)  # changed inplace op
 
         if rev:
             y2 = (x2 - t1) * torch.exp(-s1)