From 2f019604ffce3141fd0c3354ebdd1236a73a596e Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Mon, 13 Jan 2020 12:06:22 +0100
Subject: [PATCH] flow's input size is target_dim

---
 pts/model/tempflow/tempflow_network.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py
index ac5467e..9e30b68 100644
--- a/pts/model/tempflow/tempflow_network.py
+++ b/pts/model/tempflow/tempflow_network.py
@@ -54,7 +54,7 @@ class TempFlowTrainingNetwork(nn.Module):
 
         flow_cls = {"RealNVP": RealNVP}[flow_type]
         self.flow = flow_cls(
-            input_size=input_size,
+            input_size=target_dim,
             n_blocks=n_blocks,
             n_hidden=n_hidden,
             hidden_size=hidden_size,