diff --git a/pts/modules/flows.py b/pts/modules/flows.py
index b74cf8d..3522a43 100644
--- a/pts/modules/flows.py
+++ b/pts/modules/flows.py
@@ -161,12 +161,16 @@ class LinearMaskedCoupling(nn.Module):
 
         # run through model
         s = self.s_net(mx if y is None else torch.cat([y, mx], dim=-1))
-        t = self.t_net(mx if y is None else torch.cat([y, mx], dim=-1))
+        t = self.t_net(mx if y is None else torch.cat([y, mx], dim=-1)) * (1 - self.mask)
+        
         # cf RealNVP eq 8 where u corresponds to x (here we're modeling u)
-        u = mx + (1 - self.mask) * (x - t) * torch.exp(-s)
+        log_s = torch.tanh(s) * (1 - self.mask)
+        u = (x - t) * torch.exp(-log_s)
+        # u = mx + (1 - self.mask) * (x - t) * torch.exp(-s)
 
         # log det du/dx; cf RealNVP 8 and 6; note, sum over input_size done at model log_prob
-        log_abs_det_jacobian = -(1 - self.mask) * s
+        #log_abs_det_jacobian = -(1 - self.mask) * s
+        log_abs_det_jacobian = -log_s #.sum(-1, keepdim=True)
 
         return u, log_abs_det_jacobian
 
@@ -176,10 +180,14 @@ class LinearMaskedCoupling(nn.Module):
 
         # run through model
         s = self.s_net(mu if y is None else torch.cat([y, mu], dim=-1))
-        t = self.t_net(mu if y is None else torch.cat([y, mu], dim=-1))
-        x = mu + (1 - self.mask) * (u * s.exp() + t)  # cf RealNVP eq 7
+        t = self.t_net(mu if y is None else torch.cat([y, mu], dim=-1)) * (1 - self.mask)
 
-        log_abs_det_jacobian = (1 - self.mask) * s  # log det dx/du
+        log_s = torch.tanh(s) * (1 - self.mask)
+        x = u * torch.exp(log_s) + t
+        #x = mu + (1 - self.mask) * (u * s.exp() + t)  # cf RealNVP eq 7
+
+        #log_abs_det_jacobian = (1 - self.mask) * s  # log det dx/du
+        log_abs_det_jacobian = log_s #.sum(-1, keepdim=True)
 
         return x, log_abs_det_jacobian