From e0bf0bbe6a340c5f04eb91553d739e806e245e97 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Tue, 3 May 2022 17:13:04 -0700
Subject: [PATCH] Disable gradients on samplers

---
 mattport/nerf/sampler.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py
index c4b27af97..203c4639e 100644
--- a/mattport/nerf/sampler.py
+++ b/mattport/nerf/sampler.py
@@ -28,6 +28,7 @@ class UniformSampler(nn.Module):
         self.num_samples = num_samples
         self.train_stratified = train_stratified
 
+    @torch.no_grad()
     def forward(
         self,
         ray_bundle: RayBundle,
@@ -86,6 +87,7 @@ class PDFSampler(nn.Module):
         self.include_original = include_original
         self.train_stratified = train_stratified
 
+    @torch.no_grad()
     def forward(
         self,
         coarse_ray_samples: RaySamples,
@@ -106,8 +108,6 @@ class PDFSampler(nn.Module):
         Returns:
             RaySamples: Positions and deltas for samples along a ray
         """
-        # TODO (matt) Look into torch.no_grad(), or torch.inference_mode
-
         num_samples = num_samples or self.num_samples
         weights = weights[..., :-1]