[bugfix] torch profiler bug for single gpu with GPUExecutor (#8354)

2024-09-12 21:30:00 -07:00
parent 6821020109
commit ba77527955
3 changed files with 27 additions and 5 deletions
--- a/examples/offline_inference_with_profiler.py
+++ b/examples/offline_inference_with_profiler.py
@@ -16,7 +16,7 @@ prompts = [
 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

 # Create an LLM.
-llm = LLM(model="facebook/opt-125m")
+llm = LLM(model="facebook/opt-125m", tensor_parallel_size=1)

 llm.start_profile()