From 721adece9ff4ac0e020a33a1303267029705025c Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Wed, 29 Jan 2020 15:14:25 +0100 Subject: [PATCH] permute back the transformer output for conditioning --- pts/model/transformer_tempflow/transformer_tempflow_network.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pts/model/transformer_tempflow/transformer_tempflow_network.py b/pts/model/transformer_tempflow/transformer_tempflow_network.py index 7c8ac5a..1476588 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(decoder_output=dec_output) + distr_args = self.distr_args(decoder_output=dec_output.permute(1,0,2)) # (batch_size, 1, target_dim) new_samples = self.flow.sample(cond=distr_args)