diff --git a/pts/model/transformer/transformer_network.py b/pts/model/transformer/transformer_network.py
index 62262fe..c847d4d 100644
--- a/pts/model/transformer/transformer_network.py
+++ b/pts/model/transformer/transformer_network.py
@@ -83,7 +83,7 @@ class TransformerNetwork(nn.Module):
 
         # mask
         self.register_buffer(
-            "tgt_mask", torch.triu(torch.ones((prediction_length, prediction_length)))
+            "tgt_mask", torch.triu(torch.ones((prediction_length, prediction_length)), diagonal=1)
         )
 
     @staticmethod