Fix bug in copy_tree

PiperOrigin-RevId: 595951242 Change-Id: I7913d9037a5295494ca62bf314ca2e3dd4e8a008
google-deepmind · Apr 15, 2024 · 8a09e34 · 8a09e34
1 parent e94435b
commit 8a09e34
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ferminet/train.py b/ferminet/train.py
@@ -325,7 +325,8 @@ def make_kfac_training_step(
   # Due to some KFAC cleverness related to donated buffers, need to do this
   # to make state resettable
   copy_tree = constants.pmap(
-      functools.partial(jax.tree_util.tree_map, lambda x: 1.0 * x))
+      functools.partial(jax.tree_util.tree_map,
+                        lambda x: (1.0 * x).astype(x.dtype)))
 
   def step(
       data: networks.FermiNetData,