Skip to content

Commit

Permalink
Automatically identify cuda from nvidia-smi in install-requirements (p…
Browse files Browse the repository at this point in the history
…ytorch#606)

* Automatically identify cuda from nvidia-smi in install-requirements

* Update README.md

---------

Co-authored-by: Michael Gschwind <61328285+mikekgfb@users.noreply.github.com>
  • Loading branch information
2 people authored and malfet committed Jul 17, 2024
1 parent 785e92a commit 46a6b9c
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 8 deletions.
4 changes: 0 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,7 @@ cd torchchat
# set up a virtual environment
python3 -m venv .venv
source .venv/bin/activate
```

After the virtual environment is activated, dependencies can be installed.

```bash
# install dependencies
./install_requirements.sh
```
Expand Down
12 changes: 8 additions & 4 deletions install_requirements.sh
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,14 @@ $PIP_EXECUTABLE install -r requirements.txt --extra-index-url https://download.p
NIGHTLY_VERSION=dev20240422

# The pip repository that hosts nightly torch packages. cpu by default.
TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cpu"

# If cuda is available, use the pytorch nightly with cuda for faster execution on cuda GPUs.
test -f /usr/bin/nvidia-smi && TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cu121"
# If cuda is available, based on presence of nvidia-smi, install the pytorch nightly
# with cuda for faster execution on cuda GPUs.
if [[ -x "$(command -v nvidia-smi)" ]];
then
TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cu121"
else
TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/cpu"
fi

# pip packages needed by exir.
REQUIREMENTS_TO_INSTALL=(
Expand Down

0 comments on commit 46a6b9c

Please sign in to comment.