???

Author: jyc

src/levanter/models/gpt2.py

@@ -347,6 +347,9 @@ 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)
+
         input_embeds = self.token_embeddings(input_ids)
         input_embeds_norm = jnp.linalg.norm(input_embeds.array)
         input_embeds_has_nan = jnp.any(jnp.isnan(input_embeds.array))