From 7bede84f4598ae9573e7e796592226ed15367663 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Tue, 28 Jan 2020 13:44:17 +0100
Subject: [PATCH] enc_out needs to be repeated at dim=1

---
 pts/model/transformer/transformer_network.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pts/model/transformer/transformer_network.py b/pts/model/transformer/transformer_network.py
index ff92af5..62262fe 100644
--- a/pts/model/transformer/transformer_network.py
+++ b/pts/model/transformer/transformer_network.py
@@ -344,8 +344,8 @@ class TransformerPredictionNetwork(TransformerNetwork):
         ).unsqueeze(1)
 
         repeated_enc_out = enc_out.repeat_interleave(
-            repeats=self.num_parallel_samples, dim=0
-        ).unsqueeze(1)
+            repeats=self.num_parallel_samples, dim=1
+        )
 
         repeated_scale = scale.repeat_interleave(
             repeats=self.num_parallel_samples, dim=0