fix: correct error message (#193)

* fix: correct error message It was supposed to print the received array but it was not doing it. Also, the type is now displayed. * review: fix format with pyink
AI-Hypercomputer · Dec 13, 2024 · 3e46b8e · 3e46b8e
1 parent 26908e9
commit 3e46b8e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/jetstream_pt/cli.py b/jetstream_pt/cli.py
@@ -40,6 +40,7 @@
 flags.DEFINE_bool("enable_model_warmup", False, "enable model warmup")
 
 
+
 def shard_weights(env, weights, weight_shardings):
   """Shard weights according to weight_shardings"""
   sharded = {}

diff --git a/jetstream_pt/engine.py b/jetstream_pt/engine.py
@@ -310,7 +310,7 @@ def prefill(
     else:
       raise TypeError(
           "Input tokens should be of type Jax Array, but receiving:"
-          " {prefill_inputs}"
+          f" {prefill_inputs} of type {type(prefill_inputs)}"
       )
     seq_len = padded_tokens.shape[0]
     input_indexes = jnp.arange(0, seq_len)