black format

zpyoung · Jul 30, 2022 · 4ef2081 · 4ef2081
1 parent 23f7d31
commit 4ef2081
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/train_agent.py b/train_agent.py
@@ -165,7 +165,9 @@ def collect_self_play_data(
                 num_simulations_per_move,
             )
             batch = jax.device_get(batch)
-            batch = jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), batch)
+            batch = jax.tree_util.tree_map(
+                lambda x: x.reshape((-1, *x.shape[2:])), batch
+            )
             data.extend(prepare_training_data(batch, env=env))
     return data