Adding out_sharding to test model to fix recent test failure.

PiperOrigin-RevId: 718429490
google-deepmind · Jan 22, 2025 · d4388bc · d4388bc
1 parent 9182734
commit d4388bc
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/models.py b/tests/models.py
@@ -95,6 +95,7 @@ def __call__(self, inputs: LayerInputs, *_) -> LayerInputs:  # pytype: disable=s
           precision=(jax.lax.Precision.HIGHEST, jax.lax.Precision.HIGHEST),
           preferred_element_type=preferred_element_type,
           out_type=None,
+          out_sharding=None,
       )
     layer_values.append((x, y))