Move jax cache in server

c0d2734e · sbl1996@126.com · 3a79013f · c0d2734e
Commit c0d2734e authored Jul 15, 2024 by sbl1996@126.com
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

ygoinf/ygoinf/jax_inf.py ygoinf/ygoinf/jax_inf.py +4 -0

No files found.
--- a/ygoinf/ygoinf/jax_inf.py
+++ b/ygoinf/ygoinf/jax_inf.py
+import os
 import numpy as np

 import jax
+from jax.experimental.compilation_cache import compilation_cache as cc
+cc.set_cache_dir(os.path.expanduser("~/.cache/jax"))
+
 import jax.numpy as jnp
 import flax
 from ygoai.rl.jax.agent import RNNAgent