commit 0736098c9e1f9f9d580e3c1c87dcc13a56aee1b9
author Austin Schuh <austin.linux@gmail.com> Sat Oct 26 16:18:46 2024 -0700
committer Austin Schuh <austin.linux@gmail.com> Sat Oct 26 17:46:12 2024 -0700
tree aa6ab088f87d42f22ebb2f876a35ea5628616424
parent 0b0549eeac74640ebe1cab459240eae4461b3561

Quiet down printing while training

All the debugging makes it hard to see what is happening.

Change-Id: I67bba4e86f36a2dab6b780d9a85c652df141619f
Signed-off-by: Austin Schuh <austin.linux@gmail.com>

frc971/control_loops/swerve/velocity_controller/train.py

diff --git a/frc971/control_loops/swerve/velocity_controller/train.py b/frc971/control_loops/swerve/velocity_controller/train.py
index 618ab28..721f18c 100644
--- a/frc971/control_loops/swerve/velocity_controller/train.py
+++ b/frc971/control_loops/swerve/velocity_controller/train.py
@@ -424,8 +424,7 @@
     )
     state = restore_checkpoint(state, workdir)
 
-    state_sharding = nn.get_sharding(state, state.mesh)
-    logging.info(state_sharding)
+    logging.debug(nn.get_sharding(state, state.mesh))
 
     replay_buffer_state = state.replay_buffer.init({
         'observations1':
@@ -440,9 +439,7 @@
         jax.numpy.zeros((problem.num_states, )),
     })
 
-    replay_buffer_state_sharding = nn.get_sharding(replay_buffer_state,
-                                                   state.mesh)
-    logging.info(replay_buffer_state_sharding)
+    logging.debug(nn.get_sharding(replay_buffer_state, state.mesh))
 
     # Number of gradients to accumulate before doing decent.
     update_after = FLAGS.batch_size // FLAGS.num_agents