Skip to content

Y
ygo-agent
Commit c0d2734e authored by sbl1996@126.com's avatar sbl1996@126.com
Browse files
Options

Move jax cache in server

parent 3a79013f
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 0 deletions
ygoinf/ygoinf/jax_inf.py
View file @ c0d2734e
import os
import numpy as np import numpy as np
import jax 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 jax.numpy as jnp
import flax import flax
from ygoai.rl.jax.agent import RNNAgent from ygoai.rl.jax.agent import RNNAgent
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment