Fix attention_dtype flag for paged_llm_v1

Signed-off-by: aviator19941 <avinash.sharma@amd.com>
nod-ai · Feb 13, 2025 · ae89b55 · ae89b55
1 parent de70094
commit ae89b55
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/sharktank/examples/paged_llm_v1.py b/sharktank/sharktank/examples/paged_llm_v1.py
@@ -288,7 +288,7 @@ def main():
         block_seq_stride=args.block_seq_stride,
         device=device,
         activation_dtype=args.activation_dtype,
-        attention_dtype=args.activation_dtype,
+        attention_dtype=args.attention_dtype,
         attention_kernel=args.attention_kernel,
         use_hf=args.use_hf,
         tensor_parallelism_size=args.tensor_parallelism_size,