From 2d04c4dd5b002b8b988a6226d4cf7ce07162d29d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <2225664821@qq.com>
Date: Sun, 20 Aug 2023 16:35:50 +0800
Subject: [PATCH] Update modules.py

---
 modules.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/modules.py b/modules.py
index cc0aecf..92e0f32 100644
--- a/modules.py
+++ b/modules.py
@@ -370,7 +370,22 @@ class ConvFlow(nn.Module):
     unnormalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
     unnormalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
     unnormalized_derivatives = h[..., 2 * self.num_bins:]
-    
+
+    x1, logabsdet = piecewise_rational_quadratic_transform(x1,
+        unnormalized_widths,
+        unnormalized_heights,
+        unnormalized_derivatives,
+        inverse=reverse,
+        tails='linear',
+        tail_bound=self.tail_bound
+    )
+
+    x = torch.cat([x0, x1], 1) * x_mask
+    logdet = torch.sum(logabsdet * x_mask, [1,2])
+    if not reverse:
+        return x, logdet
+    else:
+        return x
 class TransformerCouplingLayer(nn.Module):
   def __init__(self,
       channels,