Return sharding whenever the leaf has .sharding attribute, not checki…

…ng explicitly for jax.Array. PiperOrigin-RevId: 565770203
google · Sep 15, 2023 · ddb04bd · ddb04bd
1 parent 654ae1a
commit ddb04bd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/training/orbax_utils.py b/flax/training/orbax_utils.py
@@ -70,7 +70,7 @@ def restore_args_from_target(target: Any, mesh: Optional[Mesh] = None) -> Any:
   """
 
   def find_sharding(x):
-    if isinstance(x, jax.Array):
+    if hasattr(x, 'sharding'):
       return x.sharding
     return None