Fixed the issue of abnormal output when the input length of timerxl is not a multiple of 96 (#15495)

jtmer · OswinGuai · web-flow · commit 25c933f53fc8 · 2025-05-14T00:31:29.000+08:00
Co-authored-by: OswinGuai &lt;peizhyi@gmail.com&gt;
diff --git a/iotdb-core/ainode/ainode/TimerXL/models/timer_xl.py b/iotdb-core/ainode/ainode/TimerXL/models/timer_xl.py
@@ -253,6 +253,18 @@ def forward(self, x, max_new_tokens: int = 96):
         self.eval()
         self.device = next(self.model.parameters()).device
         
+        if len(x.shape) == 2:
+            batch_size, cur_len = x.shape
+            if cur_len < self.config.input_token_len:
+                raise ValueError(
+                    f"Input length must be at least {self.config.input_token_len}")
+            elif cur_len % self.config.input_token_len != 0:
+                new_len = (cur_len // self.config.input_token_len) * \
+                    self.config.input_token_len
+                x = x[:, -new_len:]
+        else:
+            raise ValueError('Input shape must be: [batch_size, seq_len]')
+        
         use_cache = self.config.use_cache
         all_input_ids = x
         
