From d05fc7a2422036aed6f5e65dfe4c803057103877 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Wed, 29 Jan 2020 14:42:37 +0100
Subject: [PATCH] added missing tgt_mask

---
 .../transformer_tempflow/transformer_tempflow_network.py     | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pts/model/transformer_tempflow/transformer_tempflow_network.py b/pts/model/transformer_tempflow/transformer_tempflow_network.py
index 6401065..55f7076 100644
--- a/pts/model/transformer_tempflow/transformer_tempflow_network.py
+++ b/pts/model/transformer_tempflow/transformer_tempflow_network.py
@@ -84,6 +84,11 @@ class TransformerTempFlowTrainingNetwork(nn.Module):
 
         self.scaler = MeanScaler(keepdim=True)
 
+        # mask
+        self.register_buffer(
+            "tgt_mask", self.transformer.generate_square_subsequent_mask(prediction_length)
+        )
+
     @staticmethod
     def get_lagged_subsequences(
         sequence: torch.Tensor,