Fix sharding arg

Signed-off-by: aviator19941 <avinash.sharma@amd.com>
nod-ai · Oct 31, 2024 · d683163 · d683163
1 parent 87c4bfd
commit d683163
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/sharktank/sharktank/examples/sharding/shard_llm_dataset.py b/sharktank/sharktank/examples/sharding/shard_llm_dataset.py
@@ -30,7 +30,7 @@ def main(raw_args=None):
     args = cli.parse(parser, args=raw_args)
     dataset = cli.get_input_dataset(args)
 
-    if args.output_file is None:
+    if args.output_irpa_file is None:
         raise RuntimeError(f"Need file destination for IRPA file")
 
     if args.tensor_parallelism_size < 2:

diff --git a/sharktank/sharktank/utils/export_artifacts.py b/sharktank/sharktank/utils/export_artifacts.py
@@ -132,7 +132,7 @@ def shard_irpa_file(
             "sharktank.examples.sharding.shard_llm_dataset",
             "--gguf-file",
             gguf_file,
-            "--output-irpa",
+            "--output-irpa-file",
             output_irpa,
             "--tensor-parallelism-size",
             str(self.tensor_parallelism_size),