add name

tensorflow · Jul 20, 2021 · 4dbc208 · 4dbc208
1 parent 64b70b4
commit 4dbc208
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorflow_addons/optimizers/gradient_accumulator.py b/tensorflow_addons/optimizers/gradient_accumulator.py
@@ -57,8 +57,8 @@ def _accum_grad(grads_and_vars):
             with tf.init_scope():
                 if not self._gradients:
                     for grad, var in grads_and_vars:
-                        self._gradients[var.ref()] = tf.Variable(
-                            tf.zeros_like(var), trainable=False
+                        self._gradients[var.ref()] = self.add_weight(
+                            "ga", shape=var.shape, dtype=var.dtype, trainable=False
                         )
             new_grads_and_vars = []
             for grad, var in grads_and_vars: