From 2f93d97fd02e9d096179fb6c4215b2614c3ce42a Mon Sep 17 00:00:00 2001
From: Abhishek Jindal <abjindal@microsoft.com>
Date: Fri, 8 Dec 2023 23:12:48 -0800
Subject: [PATCH] Add cuda visible devices for Mistral benchmark (#18764)

### Description
<!-- Describe your changes. -->
Add cuda visible devices for Mistral benchmark as it is not working for
Torch compile and throwing an error.


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
Error:
File
"/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/_inductor/triton_heuristics.py",
line 556, in run
    return launcher(
  File "<string>", line 8, in launcher
RuntimeError: Triton Error [CUDA]: invalid device context
---
 .../python/tools/transformers/models/llama/README.md        | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/python/tools/transformers/models/llama/README.md b/onnxruntime/python/tools/transformers/models/llama/README.md
index 0e34fb0e69..e7bcc19635 100644
--- a/onnxruntime/python/tools/transformers/models/llama/README.md
+++ b/onnxruntime/python/tools/transformers/models/llama/README.md
@@ -412,7 +412,7 @@ python -m models.llama.convert_to_onnx -i /path/to/model/directory -o /path/to/o
 The benchmarking scripts in the LLaMA directory support Mistral benchmarking. To benchmark the ORT version, you can run: 
 
 ```
-python -m models.llama.benchmark \
+CUDA_VISIBLE_DEVICES=0 python -m models.llama.benchmark \
     -bt ort-convert-to-onnx \
     -p fp16 \
     -m mistralai/Mistral-7B-v0.1 \
@@ -422,7 +422,7 @@ python -m models.llama.benchmark \
 To benchmark the Hugging Face implementation without `torch.compile`:
 
 ```
-python -m models.llama.benchmark \
+CUDA_VISIBLE_DEVICES=0 python -m models.llama.benchmark \
     -bt hf-pt-eager \
     -p fp16 \
     -m mistralai/Mistral-7B-v0.1
@@ -431,7 +431,7 @@ python -m models.llama.benchmark \
 And to benchmark the Hugging Face implementation with `torch.compile`:
 
 ```
-python -m models.llama.benchmark \
+CUDA_VISIBLE_DEVICES=0 python -m models.llama.benchmark \
     -bt hf-pt-compile \
     -p fp16 \
     -m mistralai/Mistral-7B-v0.1