[imitation] Prevent overflow in total_frames.

vladfi1 · Jan 20, 2025 · 598b331 · 598b331
1 parent 36dca31
commit 598b331
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/slippi_ai/train_lib.py b/slippi_ai/train_lib.py
@@ -298,7 +298,8 @@ def train(config: Config):
   train_stats, _ = train_manager.step()
   logging.info('loss initial: %f', _get_loss(train_stats))
 
-  step = tf.Variable(0, trainable=False, name="step")
+  with tf.device('/cpu:0'):
+    step = tf.Variable(0, trainable=False, name="step", dtype=tf.int64)
 
   # saving and restoring
   tf_state = dict(
@@ -400,7 +401,7 @@ def maybe_log(train_stats: dict):
     print()
 
   def maybe_eval():
-    total_steps = step.numpy()
+    total_steps = int(step.numpy())
     if total_steps % runtime.eval_every_n != 0:
       return