From 9c5423f18170e0727146ae08398afce7f085f13e Mon Sep 17 00:00:00 2001
From: machenmusik <machenmusik@users.noreply.github.com>
Date: Tue, 14 Feb 2023 08:59:26 -0500
Subject: [PATCH] Update to allow pytorch (and functorch) 1.13.1 (#1390)

* Update pytorch (and functorch) to 1.13.1

Note that for existing installations, the following is recommended due to complications with older versions of functorch:
- `pip uninstall torch torchvision functorch`
- install the correct pytorch 1.13.1 version for your CUDA installation from https://pytorch.org/get-started/locally/ but add functorch to the command -- for example, Windows CUDA 11.7 would be `pip3 install torch torchvision functorch--extra-index-url https://download.pytorch.org/whl/cu117`
- reinstall nerfstudio e.g. `pip install -e .` from the source directory

* Update installation.md

* Allow 1.13.1, don't force

* pytorch 1.12.1 and 1.13.1

* Update README.md

* Update installation.md
---
 README.md                       | 14 ++++++++++++-
 docs/quickstart/installation.md | 35 ++++++++++++++++++++++++++++-----
 pyproject.toml                  |  4 ++--
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 1241b2187..7ea21deb8 100644
--- a/README.md
+++ b/README.md
@@ -95,13 +95,25 @@ python -m pip install --upgrade pip
 
 ### Dependencies
 
-Install pytorch with CUDA (this repo has been tested with CUDA 11.3) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+Install pytorch with CUDA (this repo has been tested with CUDA 11.3 and CUDA 11.7) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+
+For CUDA 11.3:
 
 ```bash
 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
 pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
 ```
 
+For CUDA 11.7:
+
+```bash
+pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
+pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+```
+
+See [Dependencies](https://github.com/nerfstudio-project/nerfstudio/blob/main/docs/quickstart/installation.md#dependencies)
+in the Installation documentation for more.
+
 ### Installing nerfstudio
 
 Easy option:
diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md
index 97c63aac8..20ab8859c 100644
--- a/docs/quickstart/installation.md
+++ b/docs/quickstart/installation.md
@@ -17,12 +17,38 @@ python -m pip install --upgrade pip
 
 ## Dependencies
 
-Install pytorch with CUDA (this repo has been tested with CUDA 11.3) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+(pytorch)=
+### pytorch
+
+Either pytorch 1.12.1 (with CUDA up to 11.3) or 1.13.1 (with CUDA 11.6/7/8) can be used.
+
+- To install 1.12.1 with CUDA 11.3:
 
 ```bash
 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
-pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+```
 
+- To install 1.13.1 with CUDA 11.7:
+
+Note that if a pytorch version prior to 1.13 is installed, 
+it should be uninstalled first to avoid upgrade issues (e.g. with functorch)
+
+```bash
+pip uninstall torch torchvision functorch
+```
+
+Install pytorch 1.13.1 with CUDA (this repo has been tested with CUDA 11.7.1) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+
+```bash
+pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
+```
+
+### tinycudann
+
+After pytorch, install the torch bindings for tinycudann:
+
+```bash
+pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
 ```
 
 ## Installing nerfstudio
@@ -134,9 +160,8 @@ While installing tiny-cuda, you run into: `The detected CUDA version mismatches
 
 **Solution**:
 
-```
-pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
-```
+Reinstall pytorch with the correct CUDA version.
+See [pytorch](pytorch) under Dependencies, above.
 
  <br />
 
diff --git a/pyproject.toml b/pyproject.toml
index 4f72b094f..8901aa282 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,7 @@ dependencies = [
     "tyro>=0.3.31",
     "gdown>=4.6.0",
     "ninja>=1.10",
-    "functorch==0.2.1",
+    "functorch>=0.2.1",
     "h5py>=2.9.0",
     "imageio>=2.21.1",
     'importlib-metadata>=6.0.0; python_version < "3.10"',
@@ -45,7 +45,7 @@ dependencies = [
     "rich>=12.5.1",
     "scikit-image>=0.19.3",
     "tensorboard==2.9.0",
-    "torch==1.12.1",
+    "torch>=1.12.1",
     "torchmetrics[image]>=0.9.3",
     "torchtyping>=0.1.4",
     "torchvision>=0.13.0",