From 9c4f7d44562d494c53dbb7bb2fb54bde5e2f56fc Mon Sep 17 00:00:00 2001
From: Mattie Tesfaldet <mattie@meta.com>
Date: Fri, 7 Jul 2023 16:25:39 -0700
Subject: [PATCH] Fixed gpu+cpu-based installation requirements in
 environment.yaml and requirements.txt.

---
 environment.yaml | 5 +++++
 requirements.txt | 4 ++++
 2 files changed, 9 insertions(+)

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