Fix parameter

ROCm · Mar 21, 2024 · ed96036 · ed96036
1 parent 97e1978
commit ed96036
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/benchmarks/benchmark_throughput.py b/benchmarks/benchmark_throughput.py
@@ -9,6 +9,7 @@
 from transformers import (AutoModelForCausalLM, AutoTokenizer,
                           PreTrainedTokenizerBase)
 from tqdm import tqdm
+from torch.distributed import launch
 
 
 def sample_requests(
@@ -93,7 +94,7 @@ def run_vllm(
               scales_path=scales_path,
               device=device,
               enable_prefix_caching=enable_prefix_caching,
-              worker_use_torchrun=args.worker_use_torchrun,)
+              worker_use_torchrun=worker_use_torchrun,)
 
     # Add the requests to the engine.
     for prompt, _, output_len in requests: