diff --git a/pts/model/transformer_tempflow/transformer_tempflow_network.py b/pts/model/transformer_tempflow/transformer_tempflow_network.py
index cd82f10..7c8ac5a 100644
--- a/pts/model/transformer_tempflow/transformer_tempflow_network.py
+++ b/pts/model/transformer_tempflow/transformer_tempflow_network.py
@@ -583,7 +583,7 @@ class TransformerTempFlowPredictionNetwork(TransformerTempFlowTrainingNetwork):
             #     unroll_length=1,
             # )
 
-            distr_args = self.distr_args(dec_output=dec_output)
+            distr_args = self.distr_args(decoder_output=dec_output)
 
             # (batch_size, 1, target_dim)
             new_samples = self.flow.sample(cond=distr_args)