Potential fix for tpu dynamic

2025-06-05 21:59:24 +02:00 · 2023-04-28 10:17:48 -05:00
parent 19817a271b
commit fa6bb4b956
1 changed files with 1 additions and 0 deletions
--- 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