Fix TVAE loss_function (#144)

* Fixes the issue

* Fix lint

Author: fealho

ctgan/synthesizers/tvae.py

@@ -54,7 +54,7 @@ def loss_function(recon_x, x, sigmas, mu, logvar, output_info, factor):
     loss = []
     for column_info in output_info:
         for span_info in column_info:
-            if len(column_info) != 1 or span_info.activation_fn != "softmax":
+            if span_info.activation_fn != "softmax":
                 ed = st + span_info.dim
                 std = sigmas[st]
                 loss.append(((x[:, st] - torch.tanh(recon_x[:, st])) ** 2 / 2 / (std ** 2)).sum())

tests/integration/test_tvae.py

@@ -51,3 +51,23 @@ def test_drop_last_false():
             correct += 1
 
     assert correct >= 95
+
+
+def test_loss_function():
+    data = pd.DataFrame({
+        '1': [float(i) for i in range(1000)],
+        '2': [float(2 * i) for i in range(1000)]
+    })
+
+    tvae = TVAESynthesizer(epochs=300)
+    tvae.fit(data)
+
+    num_samples = 1000
+    sampled = tvae.sample(num_samples)
+    error = 0
+    for _, row in sampled.iterrows():
+        error += abs(2 * row['1'] - row['2'])
+
+    avg_error = error / num_samples
+
+    assert avg_error < 400