OOM because 64gb wasn't enough 😳

markscsmith · Apr 16, 2024 · f2e5c35 · f2e5c35
1 parent e263500
commit f2e5c35
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/onefiveone/learn.py b/onefiveone/learn.py
@@ -836,8 +836,8 @@ def train_model(env, total_steps, n_steps, batch_size, episode, file_name, save_
 
 
     batch_size = 256
-    n_steps = batch_size * 32
-    total_steps = n_steps * 64
+    n_steps = batch_size * 16
+    total_steps = n_steps * 32
 
     for e in range(0, episodes):
         model = train_model(env=run_env, total_steps=total_steps, n_steps = n_steps, batch_size = batch_size, episode=e,