forked from ray-project/ray
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dl-gpu-requirements.txt
16 lines (14 loc) · 828 Bytes
/
dl-gpu-requirements.txt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# If you make changes below this line, please also make the corresponding changes to `dl-cpu-requirements.txt`!
tensorflow==2.11.0; sys_platform != 'darwin' or platform_machine != 'arm64'
tensorflow-macos==2.11.0; sys_platform == 'darwin' and platform_machine == 'arm64'
tensorflow-probability==0.19.0
tensorflow-datasets
--extra-index-url https://download.pytorch.org/whl/cu118 # for GPU versions of torch, torchvision
--find-links https://data.pyg.org/whl/torch-2.0.1+cu118.html # for GPU versions of torch-scatter, torch-sparse, torch-cluster, torch-spline-conv
# specifying explicit plus-notation below so pip overwrites the existing cpu verisons
torch==2.0.1+cu118
torchvision==0.15.2+cu118
torch-scatter==2.1.1+pt20cu118
torch-sparse==0.6.17+pt20cu118
torch-cluster==1.6.1+pt20cu118
torch-spline-conv==1.2.2+pt20cu118