Spaces:

flax-community
/

dalle-mini

Running

boris commited on Jan 14, 2022

Commit

4d518c7

1 Parent(s): db882b8

feat: add best_effort_memory_usage_reduction

Files changed (1) hide show

tools/train/train.py CHANGED Viewed

@@ -593,6 +593,7 @@ def main():
             skip_preconditioning_dim_size_gt=4096,
             clip_by_scaled_gradient_norm=None,
             precision=jax.lax.Precision.HIGHEST,
         )
     else:

             skip_preconditioning_dim_size_gt=4096,
             clip_by_scaled_gradient_norm=None,
             precision=jax.lax.Precision.HIGHEST,
+            best_effort_memory_usage_reduction=False,
         )
     else: