From fa6bb4b956a76bbe17669737c220c09a1c2c6e6c Mon Sep 17 00:00:00 2001
From: one-some <one-some+noreply@github.com>
Date: Fri, 28 Apr 2023 10:17:48 -0500
Subject: [PATCH] Potential fix for tpu dynamic

---
 tpu_mtj_backend.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tpu_mtj_backend.py b/tpu_mtj_backend.py
index d2021882..cbf24a02 100644
--- a/tpu_mtj_backend.py
+++ b/tpu_mtj_backend.py
@@ -216,6 +216,7 @@ def kobold_sample_dynamic(key, logits, rpargs, sampler_order: Optional[np.ndarra
     '''
     for sid in jnp.array(sampler_order, int):
         # sid = int(sid)
+        sid = sid.astype(int)
         warper = warpers.Warper.from_id(sid)
         if not warper.value_is_valid():
             continue