oops

Author: jyc

src/levanter/models/gpt2.py

@@ -347,8 +347,8 @@ def init(Vocab: Axis, config: Gpt2Config, *, key) -> "Gpt2Embeddings":
 
     @named_call
     def embed(self, input_ids, *, key, pos_ids: NamedArray):
-        jax.debug.print("input_ids: has_nan={nan}", nan=input_ids.has_nan)
-        jax.debug.print("token_embeddings.weight: has_nan={nan}", nan=self.token_embeddings.weight.array.has_nan)
+        jax.debug.print("input_ids: has_nan={nan}", nan=jnp.any(jnp.isnan(input_ids.array)))
+        jax.debug.print("token_embeddings.weight: has_nan={nan}", nan=jnp.any(jnp.isnan(self.token_embeddings.weight.array)))
 
         input_embeds = self.token_embeddings(input_ids)
         input_embeds_norm = jnp.linalg.norm(input_embeds.array)