From c0709fdd6900a1bb481e3c0f3a2e841c7b72c3e2 Mon Sep 17 00:00:00 2001
From: Rajat Sen <rsen91@gmail.com>
Date: Sat, 9 Nov 2024 20:28:02 +0000
Subject: [PATCH] fix median index issue

---
 src/timesfm/timesfm_jax.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/timesfm/timesfm_jax.py b/src/timesfm/timesfm_jax.py
index adfbd55..1b4427f 100644
--- a/src/timesfm/timesfm_jax.py
+++ b/src/timesfm/timesfm_jax.py
@@ -89,6 +89,7 @@ class TimesFmJax(timesfm_base.TimesFmBase):
     self._pmapped_decode = None
     self._model = None
     self._train_state = None
+    self._median_index = -1
 
   def load_from_checkpoint(
       self,