Update PTST.py

PTST.py

@@ -504,20 +504,23 @@ class PTST:
         ctx_n = (ctx - self._mu) / self._sd  # (L_eff, D)
 
         preds_n_all = []
-
-        # iterative chunking if steps > pred_len
-        while len(preds_n_all) < H:
+        n_done = 0
+        
+        while n_done < H:
             ctx_n_j = jnp.asarray(ctx_n[None, :, :])  # (1,L_eff,D)
             block_n = np.array(self._predict_block_norm(ctx_n_j)[0])  # (pred_len, D)
-
-            remaining = H - len(preds_n_all)
+        
+            remaining = H - n_done
             take = min(self.pred_len, remaining)
+        
             preds_n_all.append(block_n[:take])
-
-            # roll context: append predictions, keep last L_eff
+            n_done += take
+        
+            # roll context
             ctx_n = np.concatenate([ctx_n, block_n[:take]], axis=0)
             ctx_n = ctx_n[-self.L_eff:, :]
 
+
         preds_n = np.concatenate(preds_n_all, axis=0)  # (H, D)
         preds = preds_n * self._sd + self._mu
         return preds