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,