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,