From 0f5bb147695119c59f969d0cf0e59e32b26bdfcc Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Wed, 30 Mar 2022 15:12:18 +0200
Subject: [PATCH] input size

---
 tft/estimator.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tft/estimator.py b/tft/estimator.py
index d8d6c69..bcb4732 100644
--- a/tft/estimator.py
+++ b/tft/estimator.py
@@ -55,6 +55,7 @@ class TFTEstimator(PyTorchLightningEstimator):
         freq: str,
         prediction_length: int,
         context_length: Optional[int] = None,
+        input_size: int = 1,
         dropout: float = 0.1,
         embed_dim: int = 32,
         variable_dim: Optional[int] = None,
@@ -91,6 +92,7 @@ class TFTEstimator(PyTorchLightningEstimator):
         self.loss = loss
 
         # TFT
+        self.input_size = input_size
         self.embed_dim = embed_dim
         self.variable_dim = variable_dim or embed_dim
         self.num_heads = num_heads