From 248f81e6c2c7557e465f5ca61d17d091a6d4b970 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Wed, 30 Mar 2022 17:04:29 +0200
Subject: [PATCH] make mask optional for now

---
 tft/module.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tft/module.py b/tft/module.py
index 0626093..4ed3219 100644
--- a/tft/module.py
+++ b/tft/module.py
@@ -219,7 +219,7 @@ class TemporalFusionDecoder(nn.Module):
         return mask
 
     def forward(
-        self, x: torch.Tensor, static: torch.Tensor, mask: torch.Tensor
+        self, x: torch.Tensor, static: torch.Tensor, mask: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
         static = static.repeat((1, self.context_length + self.prediction_length, 1))