Update the optimizer test wrt the new DVariable update.

PiperOrigin-RevId: 551607364
keras-team · Jul 27, 2023 · 21c25fd · 21c25fd
1 parent ab566fd
commit 21c25fd
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/keras/dtensor/optimizers_test.py b/keras/dtensor/optimizers_test.py
@@ -89,7 +89,11 @@ def test_aggregate_gradients_noop(self):
  optimizer = adam.Adam(mesh=self.mesh)
 
  variable_init_value = tf.ones(shape=(), dtype=tf.float32)
- model_variable = dtensor.DVariable(variable_init_value, trainable=True)
+ model_variable = dtensor.DVariable(
+ variable_init_value,
+ trainable=True,
+ layout=dtensor.Layout.replicated(self.mesh, rank=0),
+ )
  grads = tf.ones_like(variable_init_value)
 
  grad_and_var = zip([grads], [model_variable])