Fix shard_llm_dataset check

Signed-off-by: aviator19941 <avinash.sharma@amd.com>
nod-ai · Nov 2, 2024 · 04af5af · 04af5af
1 parent fc51708
commit 04af5af
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sharktank/sharktank/examples/sharding/shard_llm_dataset.py b/sharktank/sharktank/examples/sharding/shard_llm_dataset.py
@@ -34,7 +34,9 @@ def main(raw_args=None):
         raise RuntimeError(f"Need file destination for IRPA file")
 
     if args.tensor_parallelism_size < 2:
-        raise RuntimeError(f"Expect sharding greater than 1 found {args.shard_count}")
+        raise RuntimeError(
+            f"Expect sharding greater than 1 found {args.tensor_parallelism_size}"
+        )
 
     hp = LlamaHParams.from_gguf_props(dataset.properties)
     llama_config = LlamaModelConfig(