fix type check (huggingface#12638)

hdqzzz · Jul 12, 2021 · f8f9a67 · f8f9a67
1 parent 2dd9440
commit f8f9a67
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/file_utils.py b/src/transformers/file_utils.py
@@ -1715,10 +1715,10 @@ def is_tensor(x):
             return True
 
     if is_flax_available():
-        import jaxlib.xla_extension as jax_xla
+        import jax.numpy as jnp
         from jax.core import Tracer
 
-        if isinstance(x, (jax_xla.DeviceArray, Tracer)):
+        if isinstance(x, (jnp.ndarray, Tracer)):
             return True
 
     return isinstance(x, np.ndarray)