diff --git a/environment.yaml b/environment.yaml index b44d4bb30..47b45b843 100644 --- a/environment.yaml +++ b/environment.yaml @@ -9,6 +9,7 @@ name: myenv channels: - pytorch + - nvidia # comment this out if using cpuonly - conda-forge - defaults @@ -20,8 +21,12 @@ channels: # users, but allows for different minor and patch versions of packages where backwards # compatibility is usually guaranteed +# Note: this is assuming a Linux-based installation, for other platforms see https://pytorch.org/get-started/locally/ + dependencies: - pytorch=2.* + - pytorch-cuda=11.8 # comment this out if using cpuonly + # - cpuonly # uncomment if you want a cpu-only installation - torchvision=0.* - lightning=2.* - torchmetrics=0.* diff --git a/requirements.txt b/requirements.txt index 668a0ed06..03e7b6e39 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,8 @@ +# Note: this is assuming a Linux-based installation, for other platforms see https://pytorch.org/get-started/locally/ + # --------- pytorch --------- # +--extra-index-url https://download.pytorch.org/whl/cu118 # gpu+cpu installation (comment this line if doing cpu-only) +# --extra-index-url https://download.pytorch.org/whl/cpu # cpu-only installation (uncomment line if doing cpu-only) torch>=2.0.0 torchvision>=0.15.0 lightning>=2.0.0