Configure jit-compilation for q-learning.

vladfi1 · Oct 20, 2024 · 8711fb3 · 8711fb3
1 parent c3fb7e3
commit 8711fb3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/slippi_ai/eval_lib.py b/slippi_ai/eval_lib.py
@@ -117,6 +117,8 @@ def multi_sample(
       return actions, hidden_state
 
     if run_on_cpu:
+      if jit_compile:
+        raise UserWarning("jit compilation may ignore run_on_cpu")
       sample = tf_utils.run_on_cpu(sample)
       multi_sample = tf_utils.run_on_cpu(multi_sample)
 

diff --git a/slippi_ai/train_q_lib.py b/slippi_ai/train_q_lib.py
@@ -49,6 +49,7 @@ class RuntimeConfig:
 class AgentConfig:
   batch_steps: int = 0
   compile: bool = True
+  jit_compile: bool = True
   name: str = nametags.DEFAULT_NAME
   async_inference: bool = False