enable bf16 casting

fattorib · Jun 4, 2023 · 668d029 · 668d029
1 parent a117722
commit 668d029
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/src/training/train_functions.py b/src/training/train_functions.py
@@ -11,11 +11,9 @@
 
 
 def to_bf16(t):
-    # return jax.tree_map(
-    #     lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t
-    # )
-    return t 
-
+    return jax.tree_map(
+        lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t
+    )