diff --git a/Dockerfile b/Dockerfile
index 92af70932..755b67c52 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -114,8 +114,8 @@ SHELL ["/bin/bash", "-c"]
 
 # Upgrade pip and install packages.
 RUN python3.10 -m pip install --upgrade pip setuptools pathtools promise pybind11
-# Install pytorch and submodules (Currently, we still use cu116 which is the latest version for torch 1.12.1 and is compatible with CUDA 11.8).
-RUN python3.10 -m pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+# Install pytorch and submodules
+RUN python3.10 -m pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
 # Install tynyCUDNN (we need to set the target architectures as environment variable first).
 ENV TCNN_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}
 RUN python3.10 -m pip install git+https://github.com/NVlabs/tiny-cuda-nn.git@v1.6#subdirectory=bindings/torch
diff --git a/README.md b/README.md
index 51a57d2a0..48bd319bc 100644
--- a/README.md
+++ b/README.md
@@ -83,7 +83,7 @@ For more complex changes (e.g., running with your own data/setting up a new NeRF
 
 ### Prerequisites
 
-You must have an NVIDIA video card with CUDA installed on the system. This library has been tested with version 11.3 of CUDA. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html)
+You must have an NVIDIA video card with CUDA installed on the system. This library has been tested with version 11.8 of CUDA. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html)
 
 ### Create environment
 
@@ -97,19 +97,19 @@ python -m pip install --upgrade pip
 
 ### Dependencies
 
-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 ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
-```
+Install pytorch with CUDA (this repo has been tested with CUDA 11.7 and CUDA 11.8) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
 
 For CUDA 11.7:
 
 ```bash
-pip install torch==1.13.1 torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
+pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
+pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+```
+
+For CUDA 11.8:
+
+```bash
+pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
 pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
 ```
 
diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md
index 58f9ffe1d..267db4776 100644
--- a/docs/quickstart/installation.md
+++ b/docs/quickstart/installation.md
@@ -4,20 +4,13 @@
 ::::::{tab-set}
 :::::{tab-item} Linux
 
-Install CUDA. This library has been tested with version 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
+Install CUDA. This library has been tested with version 11.8. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
 
 :::::
 :::::{tab-item} Windows
 Install [Git](https://git-scm.com/downloads).
 
-Install CUDA. This library and the Windows installation process has been tested with version 11.3.
-
-::::{tab-set}
-:::{tab-item} 11.3 with Visual Studio 2019 v16.9
-
-Install Visual Studio 2019 ver 16.9. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition). You can find older versions of Visual Studio 2019 [here](https://learn.microsoft.com/en-us/visualstudio/releases/2019/history#release-dates-and-build-numbers). The latest version, 16.11, may cause errors when installing tiny-cuda-nn.
-
-Install CUDA 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
+Install CUDA. This library and the Windows installation process has been tested with version 11.7.
 
 :::
 :::{tab-item} 11.7 with Visual Studio 2022
@@ -48,30 +41,39 @@ python -m pip install --upgrade pip
 ### pytorch
 
 ::::{tab-set}
-:::{tab-item} Torch 1.12.1 with CUDA 11.3
+:::{tab-item} Torch 2.0.1 with CUDA 11.8
 
-- To install 1.12.1 with CUDA 11.3:
+- To install 2.0.1 with CUDA 11.8:
 
-```bash
-pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
-```
-
-:::
-:::{tab-item} Torch 1.13.1 with CUDA 11.7
-
-- To install 1.13.1 with CUDA 11.7:
-
-Note that if a pytorch version prior to 1.13 is installed, 
+Note that if a pytorch version prior to 2.0.1 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 and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+Install pytorch 2.0.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
 
 ```bash
-pip install torch==1.13.1 torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
+pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
+```
+
+:::
+:::{tab-item} Torch 2.0.1 with CUDA 11.7
+
+- To install 2.0.1 with CUDA 11.7:
+
+Note that if a pytorch version prior to 2.0.1 is installed, 
+it should be uninstalled first to avoid upgrade issues (e.g. with functorch)
+
+```bash
+pip uninstall torch torchvision functorch
+```
+
+Install pytorch 2.0.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
+
+```bash
+pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
 ```
 
 :::
@@ -194,7 +196,6 @@ K80 | 37
 
 - [TinyCUDA installation errors out with cuda mismatch](tiny-cuda-mismatch-error)
 - [TinyCUDA installation errors out with no CUDA toolset found](tiny-cuda-integration-error)
-- [TinyCUDA installation errors out with syntax errors](tiny-cuda-syntax-error)
 - [Installation errors, File "setup.py" not found](pip-install-error)
 - [Runtime errors, "len(sources) > 0".](cuda-sources-error)
 
@@ -221,7 +222,7 @@ While installing tiny-cuda on Windows, you run into: `No CUDA toolset found.`
 
 **Solution**:
 
-Confirm that you have Visual Studio installed (CUDA 11.3 is only compatible up to 2019 ver 16.9, not 16.11 or 2022).
+Confirm that you have Visual Studio installed.
 
 Make sure CUDA Visual Studio integration is enabled. This should be done automatically by the CUDA installer if it is run after Visual Studio is installed. You can also manually enable integration.
 
@@ -231,7 +232,7 @@ Make sure CUDA Visual Studio integration is enabled. This should be done automat
 To manually enable integration for Visual Studio 2019, copy all 4 files from
 
 ```
-C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\extras\visual_studio_integration\MSBuildExtensions
+C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\extras\visual_studio_integration\MSBuildExtensions
 ```
 
 to
@@ -246,7 +247,7 @@ C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\MSBuild\Microsoft
 To manually enable integration for Visual Studio 2022, copy all 4 files from
 
 ```
-C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\extras\visual_studio_integration\MSBuildExtensions
+C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\extras\visual_studio_integration\MSBuildExtensions
 ```
 
 to
@@ -260,20 +261,6 @@ C:\Program Files\Microsoft Visual Studio\2022\[Community, Professional, Enterpri
 
  <br />
 
-(tiny-cuda-syntax-error)=
-
-**(Windows) TinyCUDA installation errors out with syntax errors**
-
-While installing tiny-cuda on Windows, you run into: `Expected a "("`
-
-**Solution**:
-
-If you are using CUDA 11.3 with Visual Studio 2019, confirm that Visual Studio is on version 16.9 and not 16.11.
-
-If you are using a different version of CUDA and Visual Studio, try switching to CUDA 11.3 with Visual Studio 2019 ver 16.9.
-
- <br />
-
 (pip-install-error)=
 
 **Installation errors, File "setup.py" not found**
diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py
index c5ac2fdf2..2c0b746d1 100644
--- a/nerfstudio/cameras/camera_utils.py
+++ b/nerfstudio/cameras/camera_utils.py
@@ -25,6 +25,8 @@ from jaxtyping import Float
 from numpy.typing import NDArray
 from torch import Tensor
 
+from nerfstudio.utils.misc import torch_compile
+
 _EPS = np.finfo(float).eps * 4.0
 
 
@@ -340,7 +342,6 @@ def get_distortion_params(
     return torch.Tensor([k1, k2, k3, k4, p1, p2])
 
 
-@torch.jit.script
 def _compute_residual_and_jacobian(
     x: torch.Tensor,
     y: torch.Tensor,
@@ -406,7 +407,7 @@ def _compute_residual_and_jacobian(
     return fx, fy, fx_x, fx_y, fy_x, fy_y
 
 
-@torch.jit.script
+@torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
 def radial_and_tangential_undistort(
     coords: torch.Tensor,
     distortion_params: torch.Tensor,
diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py
index dd9a26699..a922bc5df 100644
--- a/nerfstudio/engine/schedulers.py
+++ b/nerfstudio/engine/schedulers.py
@@ -24,6 +24,14 @@ from torch.optim import Optimizer, lr_scheduler
 
 from nerfstudio.configs.base_config import InstantiateConfig
 
+try:
+    from torch.optim.lr_scheduler import (  # pylint: disable=ungrouped-imports
+        LRScheduler,
+    )
+except ImportError:
+    # Backward compatibility for PyTorch 1.x
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+
 
 @dataclass
 class SchedulerConfig(InstantiateConfig):
@@ -43,7 +51,7 @@ class Scheduler:
         self.config = config
 
     @abstractmethod
-    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
         """Abstract method that returns a scheduler object.
 
         Args:
@@ -73,7 +81,7 @@ class MultiStepScheduler(Scheduler):
 
     config: MultiStepSchedulerConfig
 
-    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
         scheduler = lr_scheduler.MultiStepLR(
             optimizer=optimizer,
             milestones=self.config.milestones,
@@ -107,7 +115,7 @@ class ExponentialDecayScheduler(Scheduler):
 
     config: ExponentialDecaySchedulerConfig
 
-    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
         if self.config.lr_final is None:
             lr_final = lr_init
         else:
@@ -154,7 +162,7 @@ class CosineDecayScheduler(Scheduler):
 
     config: CosineDecaySchedulerConfig
 
-    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+    def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
         def func(step):
             if step < self.config.warm_up_end:
                 learning_factor = step / self.config.warm_up_end
diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py
index 97f53d17b..25011c740 100644
--- a/nerfstudio/model_components/scene_colliders.py
+++ b/nerfstudio/model_components/scene_colliders.py
@@ -22,6 +22,7 @@ import torch
 from jaxtyping import Float
 from torch import Tensor, nn
 
+from nerfstudio.utils.misc import torch_compile
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.scene_box import SceneBox
 
@@ -108,7 +109,7 @@ class AABBBoxCollider(SceneCollider):
         return ray_bundle
 
 
-@torch.jit.script
+@torch_compile(dynamic=True, mode="reduce-overhead")
 def _intersect_with_sphere(
     rays_o: torch.Tensor, rays_d: torch.Tensor, center: torch.Tensor, radius: float = 1.0, near_plane: float = 0.0
 ):
diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py
index 69a9de4d0..866d280f8 100644
--- a/nerfstudio/utils/math.py
+++ b/nerfstudio/utils/math.py
@@ -21,6 +21,8 @@ import torch
 from jaxtyping import Bool, Float
 from torch import Tensor
 
+from nerfstudio.utils.misc import torch_compile
+
 
 def components_from_spherical_harmonics(
     levels: int, directions: Float[Tensor, "*batch 3"]
@@ -197,7 +199,7 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
     return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
 
 
-@torch.jit.script
+@torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
 def intersect_aabb(
     origins: torch.Tensor,
     directions: torch.Tensor,
@@ -222,11 +224,11 @@ def intersect_aabb(
     tx_min = (aabb[:3] - origins) / directions
     tx_max = (aabb[3:] - origins) / directions
 
-    t_min = torch.min(tx_min, tx_max)
-    t_max = torch.max(tx_min, tx_max)
+    t_min = torch.stack((tx_min, tx_max)).amin(dim=0)
+    t_max = torch.stack((tx_min, tx_max)).amax(dim=0)
 
-    t_min = torch.max(t_min, dim=-1).values
-    t_max = torch.min(t_max, dim=-1).values
+    t_min = t_min.amax(dim=-1)
+    t_max = t_max.amin(dim=-1)
 
     t_min = torch.clamp(t_min, min=0, max=max_bound)
     t_max = torch.clamp(t_max, min=0, max=max_bound)
@@ -322,3 +324,4 @@ def normalized_depth_scale_and_shift(
     shift[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
 
     return scale, shift
+    return scale, shift
diff --git a/nerfstudio/utils/misc.py b/nerfstudio/utils/misc.py
index 48bd8f29c..cfa5e73db 100644
--- a/nerfstudio/utils/misc.py
+++ b/nerfstudio/utils/misc.py
@@ -155,3 +155,13 @@ def strtobool(val) -> bool:
     FMI https://stackoverflow.com/a/715468
     """
     return val.lower() in ("yes", "y", "true", "t", "on", "1")
+
+
+if hasattr(torch, "compile"):
+    # PyTorch 2.0+
+    torch_compile = torch.compile
+else:
+    # PyTorch 1.x
+    def torch_compile(*args, **kwargs):
+        """Fall back to jit script if using pytorch 1.x, args ignored"""
+        return torch.jit.script
diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py
index f64414fad..734972428 100644
--- a/nerfstudio/utils/tensor_dataclass.py
+++ b/nerfstudio/utils/tensor_dataclass.py
@@ -70,11 +70,10 @@ class TensorDataclass:
         This will 1) find the broadcasted shape and 2) broadcast all fields to this shape 3)
         set _shape to be the broadcasted shape.
         """
-        if self._field_custom_dimensions is not None:
-            for k, v in self._field_custom_dimensions.items():
-                assert (
-                    isinstance(v, int) and v > 1
-                ), f"Custom dimensions must be an integer greater than 1, since 1 is the default, received {k}: {v}"
+        for k, v in self._field_custom_dimensions.items():
+            assert (
+                isinstance(v, int) and v > 1
+            ), f"Custom dimensions must be an integer greater than 1, since 1 is the default, received {k}: {v}"
 
         # Shim to prevent pyright from narrowing `self` to DataclassInstance.
         self_dc = self
@@ -130,6 +129,7 @@ class TensorDataclass:
             if isinstance(v, torch.Tensor):
                 # If custom dimension key, then we need to
                 if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions:
+                    # pylint: disable=unsubscriptable-object
                     new_dict[k] = v.broadcast_to(
                         (
                             *batch_shape,
@@ -326,7 +326,6 @@ class TensorDataclass:
                 # This is the case when we have a custom dimensions tensor
                 elif (
                     isinstance(v, torch.Tensor)
-                    and isinstance(self._field_custom_dimensions, dict)
                     and f in self._field_custom_dimensions
                     and custom_tensor_dims_fn is not None
                 ):
diff --git a/pyproject.toml b/pyproject.toml
index cd433faf3..e9245d3d0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,6 @@ dependencies = [
     "tyro>=0.3.31",
     "gdown>=4.6.0",
     "ninja>=1.10",
-    "functorch>=0.2.1",
     "h5py>=2.9.0",
     "imageio>=2.21.1",
     'importlib-metadata>=6.0.0; python_version < "3.10"',
@@ -48,9 +47,9 @@ dependencies = [
     "rich>=12.5.1",
     "scikit-image>=0.19.3",
     "tensorboard>=2.13.0",
-    "torch>=1.12.1,<2.0.0",
+    "torch>=1.13.1",
+    "torchvision>=0.14.1",
     "torchmetrics[image]>=0.9.3",
-    "torchvision>=0.13.0",
     "typing_extensions>=4.4.0",
     "viser>=0.0.5",
     "nuscenes-devkit>=1.1.1",
@@ -115,7 +114,7 @@ ns-dev-sync-viser-message-defs = "nerfstudio.scripts.viewer.sync_viser_message_d
 [options]
 # equivalent to using --extra-index-url with pip, which is needed for specifying the CUDA version torch and torchvision
 dependency_links = [
-    "https://download.pytorch.org/whl/cu113"
+    "https://download.pytorch.org/whl/cu118"
 ]
 
 [tool.setuptools.packages.find]