Spaces:

flax-community
/

dalle-mini

Running

boris commited on Feb 2, 2022

Commit

696422e

•

1 Parent(s): 5996680

fix: distributed shampoo class

Files changed (1) hide show

tools/train/distributed_shampoo.py CHANGED Viewed

@@ -1,3 +1,5 @@
 # coding=utf-8
 # Copyright 2022 The Google Research Authors.
 #
@@ -235,7 +237,7 @@ class GraftingType(enum.IntEnum):
     RMSPROP = 3
     RMSPROP_NORMALIZED = 4
     SQRT_N = 5
-    ADAGRAD_NORMALIZED = 5
 def power_iteration(

+# file from: https://github.com/google-research/google-research/blob/master/scalable_shampoo/optax/distributed_shampoo.py
 # coding=utf-8
 # Copyright 2022 The Google Research Authors.
 #
     RMSPROP = 3
     RMSPROP_NORMALIZED = 4
     SQRT_N = 5
+    ADAGRAD_NORMALIZED = 6
 def power_iteration(