diff --git a/docs/tutorials/pipelines/creating_models.ipynb b/docs/tutorials/pipelines/creating_models.ipynb
index 5c040ceb0..09d1e8861 100644
--- a/docs/tutorials/pipelines/creating_models.ipynb
+++ b/docs/tutorials/pipelines/creating_models.ipynb
@@ -25,7 +25,7 @@
                 "\n",
                 "{ref}`fields` represents the \"space\", aka. the radiance field of the NeRF. Here, we define the field as the part of the network that takes in point samples and any other conditioning, and outputs any of the `FieldHeadNames` (`nerf/field_modules/field_heads.py`).\n",
                 "\n",
-                "All fields are composed of modules continained within the `fields/` directory (e.g. `fields/modules/` or `fields/density_fields/`). We can think of field modules as modules that actually define or interact with the field.\n",
+                "All fields are composed of modules continained within the `fields/` directory (e.g. `fields/modules/`). We can think of field modules as modules that actually define or interact with the field.\n",
                 "\n",
                 "To build a NeRF field, we therefore follow these steps:\n",
                 "\n",
@@ -201,7 +201,7 @@
     ],
     "metadata": {
         "kernelspec": {
-            "display_name": "Python 3.8.13 ('nerfactory')",
+            "display_name": "Python 3.8.12 ('nerfactory')",
             "language": "python",
             "name": "python3"
         },
@@ -215,15 +215,15 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.8.13"
+            "version": "3.8.12"
         },
         "orig_nbformat": 4,
         "vscode": {
             "interpreter": {
-                "hash": "e691f033e0f2f1b9c0a11e7e81375a1e27aec87a71fd1b1eaa545f7a5e61b3f6"
+                "hash": "34c28001ff35fb390494047002768a8182dcf55b1b11415165e62ea61557ab83"
             }
         }
     },
     "nbformat": 4,
     "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/nerfactory/cuda/__init__.py b/nerfactory/cuda/__init__.py
deleted file mode 100644
index e7affe6d7..000000000
--- a/nerfactory/cuda/__init__.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# Copyright 2022 The Plenoptix Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""CUDA Functions"""
-from typing import Callable
-
-import torch
-from torch.cuda.amp import custom_bwd, custom_fwd
-
-
-def _make_lazy_cuda(name: str) -> Callable:
-    def call_cuda(*args, **kwargs):
-        # pylint: disable=import-outside-toplevel
-        from nerfactory.cuda.backend import _C
-
-        return getattr(_C, name)(*args, **kwargs)
-
-    return call_cuda
-
-
-packbits = _make_lazy_cuda("packbits")
-ray_aabb_intersect = _make_lazy_cuda("ray_aabb_intersect")
-morton3D = _make_lazy_cuda("morton3D")
-morton3D_invert = _make_lazy_cuda("morton3D_invert")
-raymarching = _make_lazy_cuda("raymarching")
-volumetric_rendering_forward = _make_lazy_cuda("volumetric_rendering_forward")
-volumetric_rendering_backward = _make_lazy_cuda("volumetric_rendering_backward")
-occupancy_query = _make_lazy_cuda("occupancy_query")
-get_camera_indices = _make_lazy_cuda("get_camera_indices")
-
-# pylint: disable=abstract-method,arguments-differ
-class VolumeRenderer(torch.autograd.Function):
-    """CUDA Volumetirc Renderer"""
-
-    @staticmethod
-    @custom_fwd(cast_inputs=torch.float32)
-    def forward(ctx, packed_info, starts, ends, sigmas, rgbs, opacities):
-        accumulated_weight, accumulated_depth, accumulated_color, mask = volumetric_rendering_forward(
-            packed_info, starts, ends, sigmas, rgbs, opacities
-        )
-        # TODO(ruilongli): accelerate for torch.no_grad()?
-        ctx.save_for_backward(
-            accumulated_weight,
-            accumulated_depth,
-            accumulated_color,
-            packed_info,
-            starts,
-            ends,
-            sigmas,
-            rgbs,
-            opacities,
-        )
-        return accumulated_weight, accumulated_depth, accumulated_color, mask
-
-    @staticmethod
-    @custom_bwd
-    def backward(ctx, grad_weight, grad_depth, grad_color, _grad_mask):
-        (
-            accumulated_weight,
-            accumulated_depth,
-            accumulated_color,
-            packed_info,
-            starts,
-            ends,
-            sigmas,
-            rgbs,
-            opacities,
-        ) = ctx.saved_tensors
-        grad_sigmas, grad_rgbs = volumetric_rendering_backward(
-            accumulated_weight,
-            accumulated_depth,
-            accumulated_color,
-            grad_weight,
-            grad_depth,
-            grad_color,
-            packed_info,
-            starts,
-            ends,
-            sigmas,
-            rgbs,
-            opacities,
-        )
-        # corresponds to the input argument list of forward()
-        return None, None, None, grad_sigmas, grad_rgbs, None
-
-
-def unpackbits(x: torch.Tensor) -> torch.Tensor:
-    """Unpack uint8 bits back to a boolen mask.
-
-    Args:
-        x: uint8 bit tensor with shape [N]
-
-    Returns:
-        unpacked boolen tensor with shape [N * 8]
-    """
-    assert x.dtype == torch.uint8 and x.dim() == 1
-    bits = x.element_size() * 8
-    mask = 2 ** torch.arange(bits - 1, -1, -1).to(x)
-    return x.unsqueeze(-1).bitwise_and(mask).ne(0).flip(-1).bool().reshape(-1)
diff --git a/nerfactory/cuda/backend.py b/nerfactory/cuda/backend.py
deleted file mode 100644
index a6c197c50..000000000
--- a/nerfactory/cuda/backend.py
+++ /dev/null
@@ -1,55 +0,0 @@
-# Copyright 2022 The Plenoptix Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Setup cuda backend."""
-import glob
-import os
-from subprocess import DEVNULL, call
-
-from rich import console
-from torch.utils.cpp_extension import load
-
-console = console.Console()
-
-PATH = os.path.dirname(os.path.abspath(__file__))
-
-
-def cuda_toolkit_available():
-    """Check if the nvcc is avaiable on the machine."""
-    # https://github.com/idiap/fast-transformers/blob/master/setup.py
-    try:
-        call(["nvcc"], stdout=DEVNULL, stderr=DEVNULL)
-        return True
-    except FileNotFoundError:
-        return False
-
-
-if cuda_toolkit_available():
-    sources = glob.glob(os.path.join(PATH, "csrc/*.cu"))
-else:
-    sources = glob.glob(os.path.join(PATH, "csrc/*.cpp"))
-
-extra_cflags = ["-O2"]
-extra_cuda_cflags = ["-O2"]
-
-with console.status(
-    "[bold yellow]Setting up CUDA (This may take a few minutes the first time)", spinner="bouncingBall"
-):
-    _C = load(
-        name="nerfactory_extension",
-        sources=sources,
-        extra_cflags=extra_cflags,
-        extra_cuda_cflags=extra_cuda_cflags,
-    )
-
-__all__ = ["_C"]
diff --git a/nerfactory/cuda/csrc/include/helpers.h b/nerfactory/cuda/csrc/include/helpers.h
deleted file mode 100644
index 61ff8e3d5..000000000
--- a/nerfactory/cuda/csrc/include/helpers.h
+++ /dev/null
@@ -1,53 +0,0 @@
-#pragma once
-
-#ifdef __CUDACC__
-#define CUDA_HOSTDEV __host__ __device__
-#else
-#define CUDA_HOSTDEV
-#endif
-
-#include <torch/extension.h>
-
-
-inline constexpr CUDA_HOSTDEV float __SQRT3() { return 1.73205080757f; }
-
-template <typename scalar_t>
-inline CUDA_HOSTDEV void __swap(scalar_t &a, scalar_t &b)
-{
-    scalar_t c = a;
-    a = b;
-    b = c;
-}
-
-inline CUDA_HOSTDEV float __clamp(float f, float a, float b) { return fmaxf(a, fminf(f, b)); }
-
-inline CUDA_HOSTDEV int __clamp(int f, int a, int b) { return std::max(a, std::min(f, b)); }
-
-inline CUDA_HOSTDEV float __sign(float x) { return copysignf(1.0, x); }
-
-inline CUDA_HOSTDEV uint32_t __expand_bits(uint32_t v)
-{
-    v = (v * 0x00010001u) & 0xFF0000FFu;
-    v = (v * 0x00000101u) & 0x0F00F00Fu;
-    v = (v * 0x00000011u) & 0xC30C30C3u;
-    v = (v * 0x00000005u) & 0x49249249u;
-    return v;
-}
-
-inline CUDA_HOSTDEV uint32_t __morton3D(uint32_t x, uint32_t y, uint32_t z)
-{
-    uint32_t xx = __expand_bits(x);
-    uint32_t yy = __expand_bits(y);
-    uint32_t zz = __expand_bits(z);
-    return xx | (yy << 1) | (zz << 2);
-}
-
-inline CUDA_HOSTDEV uint32_t __morton3D_invert(uint32_t x)
-{
-	x = x & 0x49249249;
-	x = (x | (x >> 2)) & 0xc30c30c3;
-	x = (x | (x >> 4)) & 0x0f00f00f;
-	x = (x | (x >> 8)) & 0xff0000ff;
-	x = (x | (x >> 16)) & 0x0000ffff;
-	return x;
-}
diff --git a/nerfactory/cuda/csrc/include/helpers_cuda.h b/nerfactory/cuda/csrc/include/helpers_cuda.h
deleted file mode 100644
index 7ca474309..000000000
--- a/nerfactory/cuda/csrc/include/helpers_cuda.h
+++ /dev/null
@@ -1,18 +0,0 @@
-#pragma once
-
-#include "helpers.h"
-#include <c10/cuda/CUDAGuard.h>
-
-#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x) \
-    TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
-#define CHECK_INPUT(x) \
-    CHECK_CUDA(x);     \
-    CHECK_CONTIGUOUS(x)
-#define CUDA_GET_THREAD_ID(tid, Q)                         \
-    const int tid = blockIdx.x * blockDim.x + threadIdx.x; \
-    if (tid >= Q)                                          \
-    return
-#define CUDA_N_BLOCKS_NEEDED(Q, CUDA_N_THREADS) ((Q - 1) / CUDA_N_THREADS + 1)
-#define DEVICE_GUARD(_ten) \
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(_ten));
\ No newline at end of file
diff --git a/nerfactory/cuda/csrc/intersection.cu b/nerfactory/cuda/csrc/intersection.cu
deleted file mode 100644
index d37abd81b..000000000
--- a/nerfactory/cuda/csrc/intersection.cu
+++ /dev/null
@@ -1,118 +0,0 @@
-#include "include/helpers_cuda.h"
-
-
-template <typename scalar_t>
-inline __host__ __device__ void _ray_aabb_intersect(
-    const scalar_t* rays_o,
-    const scalar_t* rays_d,
-    const scalar_t* aabb,
-    scalar_t* near,
-    scalar_t* far
-) {
-    // aabb is [xmin, ymin, zmin, xmax, ymax, zmax]
-    scalar_t tmin = (aabb[0] - rays_o[0]) / rays_d[0];
-    scalar_t tmax = (aabb[3] - rays_o[0]) / rays_d[0];
-    if (tmin > tmax) __swap(tmin, tmax);
-
-    scalar_t tymin = (aabb[1] - rays_o[1]) / rays_d[1];
-    scalar_t tymax = (aabb[4] - rays_o[1]) / rays_d[1];
-    if (tymin > tymax) __swap(tymin, tymax);
-
-    if (tmin > tymax || tymin > tmax){
-        *near = 1e10;
-        *far = 1e10;
-        return;
-    }
-
-    if (tymin > tmin) tmin = tymin;
-    if (tymax < tmax) tmax = tymax;
-
-    scalar_t tzmin = (aabb[2] - rays_o[2]) / rays_d[2];
-    scalar_t tzmax = (aabb[5] - rays_o[2]) / rays_d[2];
-    if (tzmin > tzmax) __swap(tzmin, tzmax);
-
-    if (tmin > tzmax || tzmin > tmax){
-        *near = 1e10;
-        *far = 1e10;
-        return;
-    }
-
-    if (tzmin > tmin) tmin = tzmin;
-    if (tzmax < tmax) tmax = tzmax;
-
-    *near = tmin;
-    *far = tmax;
-    return;
-}
-
-
-template <typename scalar_t>
-__global__ void kernel_ray_aabb_intersect(
-    const int N,
-    const scalar_t* rays_o,
-    const scalar_t* rays_d,
-    const scalar_t* aabb,
-    scalar_t* t_min,
-    scalar_t* t_max
-){
-    // aabb is [xmin, ymin, zmin, xmax, ymax, zmax]
-    CUDA_GET_THREAD_ID(thread_id, N);
-
-    // locate
-    rays_o += thread_id * 3;
-    rays_d += thread_id * 3;
-    t_min += thread_id;
-    t_max += thread_id;
-    
-    _ray_aabb_intersect<scalar_t>(rays_o, rays_d, aabb, t_min, t_max);
-    
-    scalar_t zero = static_cast<scalar_t>(0.f);
-    *t_min = *t_min > zero ? *t_min : zero;
-    return;
-}
-
-/**
- * @brief Ray AABB Test
- * 
- * @param rays_o Ray origins. Tensor with shape [N, 3].
- * @param rays_d Normalized ray directions. Tensor with shape [N, 3].
- * @param aabb Scene AABB [xmin, ymin, zmin, xmax, ymax, zmax]. Tensor with shape [6].
- * @return std::vector<torch::Tensor> 
- *  Ray AABB intersection {t_min, t_max} with shape [N] respectively. Note the t_min is 
- *  clipped to minimum zero. 
- */
-std::vector<torch::Tensor> ray_aabb_intersect(
-    const torch::Tensor rays_o, const torch::Tensor rays_d, const torch::Tensor aabb
-) {
-    DEVICE_GUARD(rays_o);
-    CHECK_INPUT(rays_o);
-    CHECK_INPUT(rays_d);
-    CHECK_INPUT(aabb);
-    TORCH_CHECK(rays_o.ndimension() == 2 & rays_o.size(1) == 3)
-    TORCH_CHECK(rays_d.ndimension() == 2 & rays_d.size(1) == 3)
-    TORCH_CHECK(aabb.ndimension() == 1 & aabb.size(0) == 6)
-    
-    const int N = rays_o.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads);
-
-    torch::Tensor t_min = torch::empty({N}, rays_o.options());
-    torch::Tensor t_max = torch::empty({N}, rays_o.options());
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        rays_o.scalar_type(), "ray_aabb_intersect", 
-        ([&] {
-            kernel_ray_aabb_intersect<scalar_t><<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
-                N,
-                rays_o.data_ptr<scalar_t>(),
-                rays_d.data_ptr<scalar_t>(),
-                aabb.data_ptr<scalar_t>(),
-                t_min.data_ptr<scalar_t>(),
-                t_max.data_ptr<scalar_t>()
-            );
-        })
-    );
-
-    return {t_min, t_max};
-}
\ No newline at end of file
diff --git a/nerfactory/cuda/csrc/morton3d.cu b/nerfactory/cuda/csrc/morton3d.cu
deleted file mode 100644
index 68580faa9..000000000
--- a/nerfactory/cuda/csrc/morton3d.cu
+++ /dev/null
@@ -1,108 +0,0 @@
-#include "include/helpers_cuda.h"
-
-
-template <typename index_t>
-__global__ void kernel_morton3D(
-    const int N,
-    const index_t* coords,
-    index_t* indices
-){
-    CUDA_GET_THREAD_ID(thread_id, N);
-
-    // locate
-    coords += thread_id * 3;
-    indices += thread_id;
-
-    uint32_t index = __morton3D(
-        static_cast<uint32_t>(coords[0]), 
-        static_cast<uint32_t>(coords[1]), 
-        static_cast<uint32_t>(coords[2])
-    );
-    *indices = static_cast<index_t>(index);
-    return;
-}
-
-
-template <typename index_t>
-__global__ void kernel_morton3D_invert(
-    const int N,
-    const index_t* indices,
-    index_t* coords
-){
-    CUDA_GET_THREAD_ID(thread_id, N);
-
-    // locate
-    coords += thread_id * 3;
-    indices += thread_id;
-
-    uint32_t index = static_cast<uint32_t>(indices[0]);
-    coords[0] = static_cast<index_t>(__morton3D_invert(index >> 0));
-    coords[1] = static_cast<index_t>(__morton3D_invert(index >> 1));
-    coords[2] = static_cast<index_t>(__morton3D_invert(index >> 2));
-    return;
-}
-
-/**
- * @brief Convert 3D coords to indices in z-order
- * 
- * @param coords Integer coords with shape [N, 3]
- * @return torch::Tensor Integer indices with shape [N]
- */
-torch::Tensor morton3D(const torch::Tensor coords){
-    DEVICE_GUARD(coords);
-    CHECK_INPUT(coords);
-    // NOTE(ethan): why & and not &&?
-    TORCH_CHECK(coords.ndimension() == 2 & coords.size(1) == 3)
-
-    const int N = coords.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads);
-
-    torch::Tensor indices = torch::empty({N}, coords.options());
-
-    AT_DISPATCH_INDEX_TYPES(
-        coords.scalar_type(), "morton3D", 
-        ([&] {
-            kernel_morton3D<index_t><<<blocks, threads>>>(
-                N,
-                coords.data_ptr<index_t>(),
-                indices.data_ptr<index_t>()
-            );
-        })
-    );
-
-    return indices;
-}
-
-/**
- * @brief Convert indices in z-order to 3D coords
- * 
- * @param indices Integer indices with shape [N]
- * @return torch::Tensor Integer coords with shape [N, 3]
- */
-torch::Tensor morton3D_invert(const torch::Tensor indices){
-    DEVICE_GUARD(indices);
-    CHECK_INPUT(indices);
-    TORCH_CHECK(indices.ndimension() == 1)
-
-    const int N = indices.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads);
-
-    torch::Tensor coords = torch::empty({N, 3}, indices.options());
-
-    AT_DISPATCH_INDEX_TYPES(
-        indices.scalar_type(), "morton3D_invert", 
-        ([&] {
-            kernel_morton3D_invert<index_t><<<blocks, threads>>>(
-                N,
-                indices.data_ptr<index_t>(),
-                coords.data_ptr<index_t>()
-            );
-        })
-    );
-
-    return coords;
-}
\ No newline at end of file
diff --git a/nerfactory/cuda/csrc/nerfactory.cpp b/nerfactory/cuda/csrc/nerfactory.cpp
deleted file mode 100644
index 7e3818076..000000000
--- a/nerfactory/cuda/csrc/nerfactory.cpp
+++ /dev/null
@@ -1,90 +0,0 @@
-#include "include/helpers.h"
-
-
-torch::Tensor packbits(
-    const torch::Tensor data, const float threshold
-) {};
-
-std::vector<torch::Tensor> ray_aabb_intersect(
-    const torch::Tensor rays_o,
-    const torch::Tensor rays_d,
-    const torch::Tensor aabb
-) {};
-
-torch::Tensor morton3D(const torch::Tensor coords) {};
-
-torch::Tensor morton3D_invert(const torch::Tensor indices) {};
-
-std::vector<torch::Tensor> raymarching(
-    // rays
-    const torch::Tensor rays_o, 
-    const torch::Tensor rays_d, 
-    const torch::Tensor t_min, 
-    const torch::Tensor t_max,
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const torch::Tensor grid_bitfield, 
-    // sampling args
-    const int marching_steps,
-    const int max_total_samples,
-    const int max_ray_samples,
-    const float cone_angle,
-    const float step_scale
-) {};
-
-std::vector<torch::Tensor> volumetric_rendering_forward(
-    torch::Tensor packed_info, 
-    torch::Tensor starts, 
-    torch::Tensor ends, 
-    torch::Tensor sigmas, 
-    torch::Tensor rgbs,
-    torch::Tensor opacities
-) {};
-
-std::vector<torch::Tensor> volumetric_rendering_backward(
-    torch::Tensor accumulated_weight, 
-    torch::Tensor accumulated_depth, 
-    torch::Tensor accumulated_color, 
-    torch::Tensor grad_weight, 
-    torch::Tensor grad_depth, 
-    torch::Tensor grad_color, 
-    torch::Tensor packed_info, 
-    torch::Tensor starts, 
-    torch::Tensor ends, 
-    torch::Tensor sigmas, 
-    torch::Tensor rgbs,
-    torch::Tensor opacities
-) {};
-
-torch::Tensor get_camera_indices(
-    const int* packed_info, 
-    const int* camera_indices, 
-    torch::Tensor mapping
-) {};
-
-std::vector<torch::Tensor> occupancy_query(
-    // samples
-    const torch::Tensor positions, 
-    const torch::Tensor deltas, 
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const torch::Tensor grid_bitfield
-) {};
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
-{
-    m.def("packbits", &packbits);
-    m.def("ray_aabb_intersect", &ray_aabb_intersect);
-    m.def("morton3D", &morton3D);
-    m.def("morton3D_invert", &morton3D_invert);
-    m.def("raymarching", &raymarching);
-    m.def("volumetric_rendering_forward", &volumetric_rendering_forward);
-    m.def("volumetric_rendering_backward", &volumetric_rendering_backward);
-    m.def("occupancy_query", &occupancy_query);
-}
\ No newline at end of file
diff --git a/nerfactory/cuda/csrc/nerfactory.cu b/nerfactory/cuda/csrc/nerfactory.cu
deleted file mode 100644
index ced053728..000000000
--- a/nerfactory/cuda/csrc/nerfactory.cu
+++ /dev/null
@@ -1,91 +0,0 @@
-#include "include/helpers_cuda.h"
-
-
-torch::Tensor packbits(
-    const torch::Tensor data, const float threshold
-);
-
-std::vector<torch::Tensor> ray_aabb_intersect(
-    const torch::Tensor rays_o,
-    const torch::Tensor rays_d,
-    const torch::Tensor aabb
-);
-
-torch::Tensor morton3D(const torch::Tensor coords);
-
-torch::Tensor morton3D_invert(const torch::Tensor indices);
-
-std::vector<torch::Tensor> raymarching(
-    // rays
-    const torch::Tensor rays_o, 
-    const torch::Tensor rays_d, 
-    const torch::Tensor t_min, 
-    const torch::Tensor t_max,
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const torch::Tensor grid_bitfield, 
-    // sampling args
-    const int marching_steps,
-    const int max_total_samples,
-    const int max_ray_samples,
-    const float cone_angle,
-    const float step_scale
-);
-
-std::vector<torch::Tensor> volumetric_rendering_forward(
-    torch::Tensor packed_info, 
-    torch::Tensor starts, 
-    torch::Tensor ends, 
-    torch::Tensor sigmas, 
-    torch::Tensor rgbs,
-    torch::Tensor opacities
-);
-
-std::vector<torch::Tensor> volumetric_rendering_backward(
-    torch::Tensor accumulated_weight, 
-    torch::Tensor accumulated_depth, 
-    torch::Tensor accumulated_color, 
-    torch::Tensor grad_weight, 
-    torch::Tensor grad_depth, 
-    torch::Tensor grad_color, 
-    torch::Tensor packed_info, 
-    torch::Tensor starts, 
-    torch::Tensor ends, 
-    torch::Tensor sigmas, 
-    torch::Tensor rgbs,
-    torch::Tensor opacities
-);
-
-torch::Tensor get_camera_indices(
-    torch::Tensor packed_info, 
-    torch::Tensor camera_indices, 
-    torch::Tensor mapping
-);
-
-std::vector<torch::Tensor> occupancy_query(
-    // samples
-    const torch::Tensor positions, 
-    const torch::Tensor deltas, 
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const torch::Tensor grid_bitfield
-);
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
-{
-    m.def("packbits", &packbits);
-    m.def("ray_aabb_intersect", &ray_aabb_intersect);
-    m.def("morton3D", &morton3D);
-    m.def("morton3D_invert", &morton3D_invert);
-    m.def("raymarching", &raymarching);
-    m.def("volumetric_rendering_forward", &volumetric_rendering_forward);
-    m.def("volumetric_rendering_backward", &volumetric_rendering_backward);
-    m.def("occupancy_query", &occupancy_query);
-    m.def("get_camera_indices", &get_camera_indices);
-}
\ No newline at end of file
diff --git a/nerfactory/cuda/csrc/packbits.cu b/nerfactory/cuda/csrc/packbits.cu
deleted file mode 100644
index 3b42f104e..000000000
--- a/nerfactory/cuda/csrc/packbits.cu
+++ /dev/null
@@ -1,59 +0,0 @@
-#include "include/helpers_cuda.h"
-
-
-template <typename scalar_t>
-__global__ void kernel_packbits(
-    const int N,
-    const float threshold,
-    const scalar_t* data,
-    uint8_t* bitfield
-){
-    CUDA_GET_THREAD_ID(thread_id, N);
-
-    data += thread_id * 8;
-    bitfield += thread_id;
-
-    uint8_t bits = 0;
-    #pragma unroll 8
-    for (uint8_t i = 0; i < 8; i++) {
-        bits |= (data[i] > threshold) ? ((uint8_t)1 << i) : 0;
-    }
-    bitfield[0] = bits;
-}
-
-/**
- * @brief Pack data into uint8 bits based on threshold.
- * 
- * @param data: Tensor with any shape. Total elements must be N * 8.
- * @param threshold:
- * @return Tensor: bitfield that has shape of [N]
- */
-torch::Tensor packbits(
-    const torch::Tensor data, const float threshold
-) {
-    DEVICE_GUARD(data);
-    CHECK_INPUT(data);
-    TORCH_CHECK(data.numel() % 8 == 0)
-
-    const int N = data.numel() / 8;
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads);
-
-    torch::Tensor bitfield = torch::empty(
-        {N,}, data.options().dtype(torch::kUInt8));
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        data.scalar_type(), "packbits", 
-        ([&] {
-            kernel_packbits<scalar_t><<<blocks, threads>>>(
-                N,
-                threshold,
-                data.data_ptr<scalar_t>(),
-                bitfield.data_ptr<uint8_t>()
-            );
-        })
-    );
-
-    return bitfield;
-}
\ No newline at end of file
diff --git a/nerfactory/cuda/csrc/raymarching.cu b/nerfactory/cuda/csrc/raymarching.cu
deleted file mode 100644
index 60444b0e1..000000000
--- a/nerfactory/cuda/csrc/raymarching.cu
+++ /dev/null
@@ -1,517 +0,0 @@
-#include "include/helpers_cuda.h"
-
-
-inline __device__ float min_step_size(uint32_t max_ray_samples) { 
-    return __SQRT3() / max_ray_samples; 
-}
-
-// Maximum step size is the width of the coarsest gridsize cell.
-inline __device__ float max_step_size(uint32_t grid_cascades, uint32_t grid_size) { 
-    // Note(ruilongli): use mip=0 step size?
-    // Example:
-    // if there are 3 grid cascades and the grid size is 128 resolution
-    // then we will have 2^(num_cascades-1) = 4
-    // and sqrt(3) * 4 / 128 = 0.03125
-    return __SQRT3() * (1 << (grid_cascades-1)) / grid_size; 
-}
-
-// Perform fixed-size stepping in unit-cube scenes (like original NeRF) and exponential
-// stepping in larger scenes. 
-inline __device__ float calc_dt(float t, float cone_angle, float dt_min, float dt_max) {
-    // TODO(ruilongli): scene_scale related to cone_angle?
-    return __clamp(t * cone_angle, dt_min, dt_max);
-}
-
-inline __device__ int mip_from_pos(float x, float y, float z, uint32_t grid_cascades, float grid_center, float grid_scale) {
-    float maxval = fmaxf(fmaxf(fabsf(x - grid_center), fabsf(y - grid_center)), fabsf(z - grid_center)) / grid_scale;
-    int exponent; frexpf(maxval, &exponent);
-    return min(grid_cascades-1, max(0, exponent+1));
-}
-
-inline __device__ int mip_from_dt(
-    float x, float y, float z, uint32_t grid_cascades, float dt, int grid_size, float grid_center, float grid_scale
-) {
-    int mip = mip_from_pos(x, y, z, grid_cascades, grid_center, grid_scale);
-    dt *= 2 * grid_size;
-    if (dt<1.f) return mip; // exponent would be zero
-    int exponent; frexpf(dt, &exponent);
-    return min(grid_cascades-1, max(exponent, mip));
-}
-
-inline __device__ uint32_t grid_mip_offset(uint32_t mip, int grid_size) {
-    return (grid_size * grid_size * grid_size) * mip;
-}
-
-inline __device__ uint32_t cascaded_grid_idx_at(
-    float x, float y, float z, uint32_t mip, int grid_size, float grid_center, float grid_scale
-) {
-    // TODO(ruilongli): if the x, y, z is outside the aabb, it will be clipped into aabb!!! We should just return false
-    float mip_scale = scalbnf(1.0f, -mip) / grid_scale;
-    int ix = (int)((mip_scale * (x - grid_center) + 0.5f) * grid_size);
-    int iy = (int)((mip_scale * (y - grid_center) + 0.5f) * grid_size);
-    int iz = (int)((mip_scale * (z - grid_center) + 0.5f) * grid_size);
-    // printf("[input] x %f, y %f, z %f, mip %d, grid_size %d, grid_center %f\n", x, y, z, mip, grid_size, grid_center);
-    // printf("[output] mip_scale %f, ix %d iy %d, iz %d\n", mip_scale, ix, iy, iz);
-    uint32_t idx = __morton3D(
-        __clamp(ix, 0, grid_size-1),
-        __clamp(iy, 0, grid_size-1),
-        __clamp(iz, 0, grid_size-1)
-    );
-    return idx;
-}
-
-inline __device__ bool grid_occupied_at(
-    float x, float y, float z, 
-    const uint8_t* grid_bitfield,
-    uint32_t mip, int grid_size, float grid_center, float grid_scale
-) {
-    uint32_t idx = (
-        cascaded_grid_idx_at(x, y, z, mip, grid_size, grid_center, grid_scale)
-        + grid_mip_offset(mip, grid_size)
-    );
-    return grid_bitfield[idx/8] & (1<<(idx%8));
-}
-
-inline __device__ float distance_to_next_voxel(
-    float x, float y, float z, 
-    float dir_x, float dir_y, float dir_z, 
-    float idir_x, float idir_y, float idir_z,
-    float res
-) { // dda like step
-    // TODO: warning: expression has no effect?
-    x, y, z = res * x, res * y, res * z;
-    float tx = (floorf(x + 0.5f + 0.5f * __sign(dir_x)) - x) * idir_x;
-    float ty = (floorf(y + 0.5f + 0.5f * __sign(dir_y)) - y) * idir_y;
-    float tz = (floorf(z + 0.5f + 0.5f * __sign(dir_z)) - z) * idir_z;
-    float t = min(min(tx, ty), tz);
-
-    return fmaxf(t / res, 0.0f);
-}
-
-inline __device__ float advance_to_next_voxel(
-    float t,
-    float x, float y, float z, 
-    float dir_x, float dir_y, float dir_z, 
-    float idir_x, float idir_y, float idir_z,
-    float res, float dt_min) {
-    // Regular stepping (may be slower but matches non-empty space)
-    float t_target = t + distance_to_next_voxel(
-        x, y, z, dir_x, dir_y, dir_z, idir_x, idir_y, idir_z, res
-    );
-    do {
-        t += dt_min;
-    } while (t < t_target);
-    return t;
-}
-
-
-// Two modes: 
-// 1. "vertically" marching: a portion of rays march maximum steps. (during training)
-// 2. "horizontally" marching: all rays march together for certain steps. (during inference)
-template <typename scalar_t>
-__global__ void kernel_raymarching(
-    // rays info
-    const uint32_t n_rays,
-    const scalar_t* rays_o, // shape (n_rays, 3)
-    const scalar_t* rays_d, // shape (n_rays, 3)
-    const scalar_t* t_min, // shape (n_rays,)
-    const scalar_t* t_max, // shape (n_rays,)
-    // density grid
-    const float grid_center, // default 0.0
-    const float grid_scale, // default 3.0
-    const int grid_cascades, // default 1
-    const int grid_size, // default 128
-    const uint8_t* grid_bitfield,
-    // sampling
-    const int marching_steps, // used only for inference. -1 means disable it.
-    const int max_total_samples,
-    const int max_ray_samples, // default 1024
-    const float cone_angle, // default 0. for nerf-syn and 1/256 for large scene
-    const float step_scale, // length of the longest side of the scene bounding box
-    int* steps_counter,
-    int* rays_counter,
-    int* packed_info_out,
-    scalar_t* origins_out,
-    scalar_t* dirs_out,
-    scalar_t* starts_out,
-    scalar_t* ends_out 
-) {
-    CUDA_GET_THREAD_ID(i, n_rays);
-
-    // locate
-    rays_o += i * 3;
-    rays_d += i * 3;
-    t_min += i;
-    t_max += i;
-
-    const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2];
-    const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2];
-    const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz;
-    const float near = t_min[0], far = t_max[0];
-
-    // TODO(ruilongli): perturb `near` as in ngp_pl?
-    // TODO(ruilongli): pre-compute `dt_min`, `dt_max` as it is a constant?
-    float dt_min = min_step_size(max_ray_samples) * step_scale;
-    float dt_max = max_step_size(grid_cascades, grid_size) * grid_scale;
-
-    // if (threadIdx.x == 0) {
-    //     // these values should be equivalent with the current code
-    //     printf("step_scale value %f\n", step_scale);
-    //     printf("grid_scale value %f\n", grid_scale);
-    // }
-	
-    uint32_t ray_idx, base, marching_samples;
-    uint32_t j;
-    float t0, dt, t1, t_mid;
-    if (marching_steps < 0) {
-        // first pass to compute an accurate number of steps
-        j = 0;
-        t0 = near;
-        dt = calc_dt(t0, cone_angle, dt_min, dt_max);
-        t1 = t0 + dt;
-        t_mid = (t0 + t1) * 0.5f;
-
-        while (t_mid < far && j < max_ray_samples) {
-            // current center
-            const float x = ox + t_mid * dx;
-            const float y = oy + t_mid * dy;
-            const float z = oz + t_mid * dz;
-            uint32_t mip = mip_from_dt(x, y, z, grid_cascades, dt, grid_size, grid_center, grid_scale);
-            
-            if (grid_occupied_at(x, y, z, grid_bitfield, mip, grid_size, grid_center, grid_scale)) {
-                ++j;
-                // march to next sample
-                t0 = t1;
-                dt = calc_dt(t0, cone_angle, dt_min, dt_max);
-                t1 = t0 + dt;
-                t_mid = (t0 + t1) * 0.5f;
-            }
-            else {
-                // march to next sample
-                float res = (grid_size >> mip) * grid_scale;
-                t_mid = advance_to_next_voxel(
-                    t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, res, dt_min
-                );
-                dt = calc_dt(t_mid, cone_angle, dt_min, dt_max);
-                t0 = t_mid - dt * 0.5f;
-                t1 = t_mid + dt * 0.5f;
-            }
-        }
-        if (j == 0) return;
-
-        marching_samples = j;
-        base = atomicAdd(steps_counter, marching_samples);
-        if (base + marching_samples > max_total_samples) return;
-        ray_idx = atomicAdd(rays_counter, 1);
-
-    } else {
-        marching_samples = marching_steps;
-        base = i * marching_steps;
-        ray_idx = i;
-    }
-
-    // locate
-    origins_out += base * 3;
-    dirs_out += base * 3;
-    starts_out += base;
-    ends_out += base;
-
-    // Second round
-    j = 0;
-    t0 = near;
-    dt = calc_dt(t0, cone_angle, dt_min, dt_max);
-    t1 = t0 + dt;
-    t_mid = (t0 + t1) / 2.;
-
-    while (t_mid < far && j < marching_samples) {
-        // current center
-        const float x = ox + t_mid * dx;
-        const float y = oy + t_mid * dy;
-        const float z = oz + t_mid * dz;
-        uint32_t mip = mip_from_dt(x, y, z, grid_cascades, dt, grid_size, grid_center, grid_scale);
-        
-        if (grid_occupied_at(x, y, z, grid_bitfield, mip, grid_size, grid_center, grid_scale)) {
-            origins_out[j * 3 + 0] = ox;
-            origins_out[j * 3 + 1] = oy;
-            origins_out[j * 3 + 2] = oz;
-            dirs_out[j * 3 + 0] = dx;
-            dirs_out[j * 3 + 1] = dy;
-            dirs_out[j * 3 + 2] = dz;
-            starts_out[j] = t0;   
-            ends_out[j] = t1;     
-            ++j;
-            // march to next sample
-            t0 = t1;
-            dt = calc_dt(t0, cone_angle, dt_min, dt_max);
-            t1 = t0 + dt;
-            t_mid = (t0 + t1) * 0.5f;
-        }
-        else {
-            // march to next sample
-            float res = (grid_size >> mip) * grid_scale;
-            t_mid = advance_to_next_voxel(
-                t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, res, dt_min
-            );
-            dt = calc_dt(t_mid, cone_angle, dt_min, dt_max);
-            t0 = t_mid - dt * 0.5f;
-            t1 = t_mid + dt * 0.5f;
-		}
-	}
-
-    packed_info_out[ray_idx * 3 + 0] = i;  // ray idx in {rays_o, rays_d}
-    packed_info_out[ray_idx * 3 + 1] = base;  // point idx start.
-    packed_info_out[ray_idx * 3 + 2] = j;  // point idx shift (actual marching samples).
-
-    return;
-}
-
-
-template <typename scalar_t>
-__global__ void kernel_occupancy_query(
-    // samples info
-    const uint32_t n_samples,
-    const scalar_t* positions, 
-    const scalar_t* deltas, 
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const uint8_t* grid_bitfield,
-    // outputs
-    int32_t* mip_levels, // output mip level
-    int32_t* indices,  // output indices
-    bool* occupancies  // output occupancy
-) {
-    CUDA_GET_THREAD_ID(i, n_samples);
-
-    // locate
-    positions += i * 3;
-    deltas += i;
-    mip_levels += i;
-    indices += i;
-    occupancies += i;
-
-    // current point
-    const float x = positions[0];
-    const float y = positions[1];
-    const float z = positions[2];
-    const float dt = deltas[0];
-    
-    int32_t mip = mip_from_dt(x, y, z, grid_cascades, dt, grid_size, grid_center, grid_scale);
-    mip_levels[0] = mip;
-    
-    int32_t idx = (
-        cascaded_grid_idx_at(x, y, z, mip, grid_size, grid_center, grid_scale)
-        + grid_mip_offset(mip, grid_size)
-    );
-    indices[0] = idx;
-
-    if (grid_bitfield) {
-        bool occ = grid_bitfield[idx/8] & (1<<(idx%8));
-        occupancies[0] = occ;
-    }
-    return;
-}
-
-
-/**
- * @brief Sample points by ray marching.
- * 
- * @param rays_o Ray origins Shape of [n_rays, 3].
- * @param rays_d Normalized ray directions. Shape of [n_rays, 3].
- * @param t_min Near planes of rays. Shape of [n_rays].
- * @param t_max Far planes of rays. Shape of [n_rays].
- * @param grid_center Density grid center. TODO: support 3-dims.
- * @param grid_scale Density grid base level scale. TODO: support 3-dims.
- * @param grid_cascades Density grid levels.
- * @param grid_size Density grid resolution.
- * @param grid_bitfield Density grid uint8 bit field.
- * @param marching_steps Marching steps during inference.
- * @param max_total_samples Maximum total number of samples in this batch.
- * @param max_ray_samples Used to define the minimal step size: SQRT3() / max_ray_samples.
- * @param cone_angle 0. for nerf-synthetic and 1./256 for real scenes.
- * @param step_scale Scale up the step size by this much. Usually equals to scene scale.
- *   For example, this is often set to the length of the longest side of the scene bounding box.
- * @return std::vector<torch::Tensor> 
- * - packed_info: Stores how to index the ray samples from the returned values.
- *  Shape of [n_rays, 3]. First value is the ray index. Second value is the sample 
- *  start index in the results for this ray. Third value is the number of samples for
- *  this ray. Note for rays that have zero samples, we simply skip them so the `packed_info`
- *  has some zero padding in the end.
- * - origins: Ray origins for those samples. [max_total_samples, 3]
- * - dirs: Ray directions for those samples. [max_total_samples, 3]
- * - starts: Where the frustum-shape sample starts along a ray. [max_total_samples, 1]
- * - ends: Where the frustum-shape sample ends along a ray. [max_total_samples, 1]
- */
-std::vector<torch::Tensor> raymarching(
-    // rays
-    const torch::Tensor rays_o, 
-    const torch::Tensor rays_d, 
-    const torch::Tensor t_min, 
-    const torch::Tensor t_max,
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const torch::Tensor grid_bitfield, 
-    // sampling args
-    const int marching_steps,
-    const int max_total_samples,
-    const int max_ray_samples,
-    const float cone_angle,
-    const float step_scale
-) {
-    DEVICE_GUARD(rays_o);
-
-    CHECK_INPUT(rays_o);
-    CHECK_INPUT(rays_d);
-    CHECK_INPUT(t_min);
-    CHECK_INPUT(t_max);
-    CHECK_INPUT(grid_bitfield);
-    
-    const int n_rays = rays_o.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads);
-
-    // helper counter
-    torch::Tensor steps_counter = torch::zeros(
-        {1}, rays_o.options().dtype(torch::kInt32));
-    torch::Tensor rays_counter = torch::zeros(
-        {1}, rays_o.options().dtype(torch::kInt32));
-
-    int output_total_samples = marching_steps < 0 ? max_total_samples : marching_steps * n_rays;
-
-    // output samples
-    torch::Tensor packed_info = torch::zeros(
-        {n_rays, 3}, rays_o.options().dtype(torch::kInt32));  // ray_id, sample_id, num_samples
-    torch::Tensor origins = torch::zeros({output_total_samples, 3}, rays_o.options());
-    torch::Tensor dirs = torch::zeros({output_total_samples, 3}, rays_o.options());
-    torch::Tensor starts = torch::zeros({output_total_samples, 1}, rays_o.options());
-    torch::Tensor ends = torch::zeros({output_total_samples, 1}, rays_o.options());
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        rays_o.scalar_type(),
-        "raymarching_train",
-        ([&]
-         { kernel_raymarching<scalar_t><<<blocks, threads>>>(
-                // rays
-                n_rays,
-                rays_o.data_ptr<scalar_t>(),
-                rays_d.data_ptr<scalar_t>(),
-                t_min.data_ptr<scalar_t>(),
-                t_max.data_ptr<scalar_t>(),
-                // density grid
-                grid_center,
-                grid_scale,
-                grid_cascades,
-                grid_size,
-                grid_bitfield.data_ptr<uint8_t>(),
-                // sampling args
-                marching_steps, // used only for inference. -1 means disable it.
-                max_total_samples,
-                max_ray_samples,
-                cone_angle,
-                step_scale,
-                steps_counter.data_ptr<int>(),  // total samples.
-                rays_counter.data_ptr<int>(),  // total rays.
-                packed_info.data_ptr<int>(), 
-                origins.data_ptr<scalar_t>(),
-                dirs.data_ptr<scalar_t>(), 
-                starts.data_ptr<scalar_t>(),
-                ends.data_ptr<scalar_t>()
-            ); 
-        }));
-
-    return {packed_info, origins, dirs, starts, ends};
-}
-
-__global__ void kernel_get_camera_indices(
-    const uint32_t n_rays,
-    const int* packed_info, 
-    const long* camera_indices,
-    long* mapping
-) {
-    CUDA_GET_THREAD_ID(thread_id, n_rays);
-
-    const int i = packed_info[thread_id * 3 + 0];  // ray idx in {rays_o, rays_d}
-    const int base = packed_info[thread_id * 3 + 1];  // point idx start.
-    const int numsteps = packed_info[thread_id * 3 + 2];  // point idx shift.
-    const long camera_index = camera_indices[i];
-
-    int j = 0;
-    for (; j < numsteps; ++j) {
-        mapping[base + j] = camera_index;
-    }
-    
-    return;
-}
-
-torch::Tensor get_camera_indices(
-    const torch::Tensor packed_info, 
-    const torch::Tensor camera_indices, 
-    torch::Tensor mapping
-) {
-    const int n_rays = packed_info.size(0);
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads);
-
-    kernel_get_camera_indices<<<blocks, threads>>>(n_rays, packed_info.data_ptr<int>(), camera_indices.data_ptr<long>(), mapping.data_ptr<long>());
-    return mapping;
-}
-
-
-std::vector<torch::Tensor> occupancy_query(
-    // samples
-    const torch::Tensor positions, 
-    const torch::Tensor deltas, 
-    // density grid
-    const float grid_center,
-    const float grid_scale,
-    const int grid_cascades,
-    const int grid_size,
-    const torch::Tensor grid_bitfield
-) {
-    DEVICE_GUARD(positions);
-
-    CHECK_INPUT(positions);
-    CHECK_INPUT(deltas);
-    CHECK_INPUT(grid_bitfield);
-    
-    const int n_samples = positions.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(n_samples, threads);
-
-    // outputs
-    torch::Tensor mip_levels = torch::empty({n_samples}, positions.options().dtype(torch::kInt32));
-    torch::Tensor indices = torch::empty({n_samples}, positions.options().dtype(torch::kInt32));
-    torch::Tensor occupancies = torch::empty({n_samples}, positions.options().dtype(torch::kBool));
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        positions.scalar_type(),
-        "occupancy_query",
-        ([&]
-         { kernel_occupancy_query<scalar_t><<<blocks, threads>>>(
-                // samples
-                n_samples,
-                positions.data_ptr<scalar_t>(),
-                deltas.data_ptr<scalar_t>(),
-                // density grid
-                grid_center,
-                grid_scale,
-                grid_cascades,
-                grid_size,
-                grid_bitfield.data_ptr<uint8_t>(),
-                // outputs
-                mip_levels.data_ptr<int32_t>(),
-                indices.data_ptr<int32_t>(),
-                occupancies.data_ptr<bool>()
-            ); 
-        }));
-
-    return {mip_levels, indices, occupancies};
-}
-
diff --git a/nerfactory/cuda/csrc/volume_rendering.cu b/nerfactory/cuda/csrc/volume_rendering.cu
deleted file mode 100644
index 996dacbdd..000000000
--- a/nerfactory/cuda/csrc/volume_rendering.cu
+++ /dev/null
@@ -1,271 +0,0 @@
-#include "include/helpers_cuda.h"
-
-
-template <typename scalar_t>
-__global__ void volumetric_rendering_forward_kernel(
-    const uint32_t n_rays,
-    const int* packed_info,  // input ray & point indices.
-    const scalar_t* starts,  // input start t
-    const scalar_t* ends,  // input end t
-    const scalar_t* sigmas,  // input density after activation
-    const scalar_t* rgbs,  // input rgb after activation 
-    const scalar_t* opacities,  // input opacities for each ray
-    // should be all-zero initialized
-    scalar_t* accumulated_weight,  // output
-    scalar_t* accumulated_depth,  // output
-    scalar_t* accumulated_color,  // output
-    bool* mask  // output
-) {
-    CUDA_GET_THREAD_ID(thread_id, n_rays);
-
-    // locate
-    const int i = packed_info[thread_id * 3 + 0];  // ray idx in {rays_o, rays_d}
-    const int base = packed_info[thread_id * 3 + 1];  // point idx start.
-    const int numsteps = packed_info[thread_id * 3 + 2];  // point idx shift.
-    if (numsteps == 0) return;
-
-    starts += base;
-    ends += base;
-    sigmas += base;
-    rgbs += base * 3;
-
-    accumulated_weight += i;
-    accumulated_depth += i;
-    accumulated_color += i * 3;
-    mask += i;
-    opacities += i;
-
-    // accumulated rendering
-    scalar_t T = 1.f - opacities[0];
-    scalar_t EPSILON = 1e-4f;
-    int j = 0;
-    for (; j < numsteps; ++j) {
-        if (T < EPSILON) {
-            break;
-        }
-        const scalar_t delta = ends[j] - starts[j];
-        const scalar_t t = (ends[j] + starts[j]) * 0.5f;
-
-        const scalar_t alpha = 1.f - __expf(-sigmas[j] * delta);
-        const scalar_t weight = alpha * T;
-        accumulated_weight[0] += weight;
-        accumulated_depth[0] += weight * t;
-        accumulated_color[0] += weight * rgbs[j * 3 + 0];
-        accumulated_color[1] += weight * rgbs[j * 3 + 1];
-        accumulated_color[2] += weight * rgbs[j * 3 + 2];
-        T *= (1.f - alpha);
-    }
-    mask[0] = true;
-}
-
-
-template <typename scalar_t>
-__global__ void volumetric_rendering_backward_kernel(
-    const uint32_t n_rays,
-    const int* packed_info,  // input ray & point indices.
-    const scalar_t* starts,  // input start t
-    const scalar_t* ends,  // input end t
-    const scalar_t* sigmas,  // input density after activation
-    const scalar_t* rgbs,  // input rgb after activation 
-    const scalar_t* opacities,  // input opacities for each ray
-    const scalar_t* accumulated_weight,  // forward output
-    const scalar_t* accumulated_depth,  // forward output
-    const scalar_t* accumulated_color,  // forward output
-    const scalar_t* grad_weight,  // input
-    const scalar_t* grad_depth,  // input
-    const scalar_t* grad_color,  // input
-    scalar_t* grad_sigmas,  // output
-    scalar_t* grad_rgbs  // output
-) {
-    CUDA_GET_THREAD_ID(thread_id, n_rays);
-
-    // locate
-    const int i = packed_info[thread_id * 3 + 0];  // ray idx in {rays_o, rays_d}
-    const int base = packed_info[thread_id * 3 + 1];  // point idx start.
-    const int numsteps = packed_info[thread_id * 3 + 2];  // point idx shift.
-    if (numsteps == 0) return;
-
-    starts += base;
-    ends += base;
-    sigmas += base;
-    rgbs += base * 3;
-
-    grad_sigmas += base;
-    grad_rgbs += base * 3;
-
-    accumulated_weight += i;
-    accumulated_depth += i;
-    accumulated_color += i * 3;
-    opacities += i;
-    
-    grad_weight += i;
-    grad_depth += i;
-    grad_color += i * 3;
-    
-    // backward of accumulated rendering
-    scalar_t T = 1.f - opacities[0];
-    scalar_t EPSILON = 1e-4f;
-    int j = 0;
-    scalar_t r = 0, g = 0, b = 0, d = 0;
-    for (; j < numsteps; ++j) {
-        if (T < EPSILON) {
-            break;
-        }
-        const scalar_t delta = ends[j] - starts[j];
-        const scalar_t t = (ends[j] + starts[j]) * 0.5f;
-
-        const scalar_t alpha = 1.f - __expf(-sigmas[j] * delta);
-        const scalar_t weight = alpha * T;
-
-        r += weight * rgbs[j * 3 + 0];
-        g += weight * rgbs[j * 3 + 1];
-        b += weight * rgbs[j * 3 + 2];
-        d += weight * t;
-
-        T *= (1.f - alpha);
-
-        grad_rgbs[j * 3 + 0] = grad_color[0] * weight;
-        grad_rgbs[j * 3 + 1] = grad_color[1] * weight;
-        grad_rgbs[j * 3 + 2] = grad_color[2] * weight;
-
-        grad_sigmas[j] = delta * (
-            grad_color[0] * (T * rgbs[j * 3 + 0] - (accumulated_color[0] - r)) +
-            grad_color[1] * (T * rgbs[j * 3 + 1] - (accumulated_color[1] - g)) +
-            grad_color[2] * (T * rgbs[j * 3 + 2] - (accumulated_color[2] - b)) +
-            grad_weight[0] * (1.f - accumulated_weight[0]) +
-            grad_depth[0] * (t * T - (accumulated_depth[0] - d))
-        );
-    }
-}
-
-/**
- * @brief Volumetric Rendering: Accumulating samples in the forward pass.
- *  The inputs, excepct for `sigmas` and `rgbs`, are the outputs of our
- *  cuda ray marching function in `raymarching.cu`
- * 
- * @param packed_info Stores how to index the ray samples from the returned values.
- *  Shape of [n_rays, 3]. First value is the ray index. Second value is the sample 
- *  start index in the results for this ray. Third value is the number of samples for
- *  this ray. Note for rays that have zero samples, we simply skip them so the `packed_info`
- *  has some zero padding in the end.
- * @param starts: Where the frustum-shape sample starts along a ray. [total_samples, 1]
- * @param ends: Where the frustum-shape sample ends along a ray. [total_samples, 1]
- * @param sigmas Densities at those samples. [total_samples, 1]
- * @param rgbs RGBs at those samples. [total_samples, 3]
- * @param opacities Opacities that the rays already have. Zero for training but non-zero for eval. [total_samples, 1]
- * @return std::vector<torch::Tensor> 
- * - accumulated_weight: Ray opacity. [n_rays, 1]
- * - accumulated_depth: Ray depth. [n_rays, 1]
- * - accumulated_color: Ray color. [n_rays, 3]
- * - mask: Boolen value store if this ray has valid samples from packed_info. [n_rays]
- */
-std::vector<torch::Tensor> volumetric_rendering_forward(
-    torch::Tensor packed_info, 
-    torch::Tensor starts, 
-    torch::Tensor ends, 
-    torch::Tensor sigmas, 
-    torch::Tensor rgbs,
-    torch::Tensor opacities
-) {
-    DEVICE_GUARD(packed_info);
-    CHECK_INPUT(packed_info);
-    CHECK_INPUT(starts);
-    CHECK_INPUT(ends);
-    CHECK_INPUT(sigmas);
-    CHECK_INPUT(rgbs);
-    CHECK_INPUT(opacities);
-    TORCH_CHECK(packed_info.ndimension() == 2 & packed_info.size(1) == 3);
-    TORCH_CHECK(starts.ndimension() == 2 & starts.size(1) == 1);
-    TORCH_CHECK(ends.ndimension() == 2 & ends.size(1) == 1);
-    TORCH_CHECK(sigmas.ndimension() == 2 & sigmas.size(1) == 1);
-    TORCH_CHECK(rgbs.ndimension() == 2 & rgbs.size(1) == 3);
-    TORCH_CHECK(opacities.ndimension() == 2 & opacities.size(1) == 1);
-
-    const uint32_t n_rays = packed_info.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads);
-
-    // outputs
-    torch::Tensor accumulated_weight = torch::zeros({n_rays, 1}, sigmas.options()); 
-    torch::Tensor accumulated_depth = torch::zeros({n_rays, 1}, sigmas.options()); 
-    torch::Tensor accumulated_color = torch::zeros({n_rays, 3}, sigmas.options()); 
-    // The rays that are not skipped during sampling.
-    torch::Tensor mask = torch::zeros({n_rays}, sigmas.options().dtype(torch::kBool)); 
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        sigmas.scalar_type(),
-        "volumetric_rendering_forward",
-        ([&]
-         { volumetric_rendering_forward_kernel<scalar_t><<<blocks, threads>>>(
-                n_rays,
-                packed_info.data_ptr<int>(), 
-                starts.data_ptr<scalar_t>(),
-                ends.data_ptr<scalar_t>(),
-                sigmas.data_ptr<scalar_t>(),
-                rgbs.data_ptr<scalar_t>(),
-                opacities.data_ptr<scalar_t>(),
-                accumulated_weight.data_ptr<scalar_t>(),
-                accumulated_depth.data_ptr<scalar_t>(),
-                accumulated_color.data_ptr<scalar_t>(),
-                mask.data_ptr<bool>()
-            ); 
-        }));
-
-    return {accumulated_weight, accumulated_depth, accumulated_color, mask};
-}
-
-
-
-/**
- * @brief Volumetric Rendering: Accumulating samples in the backward pass.
- */
-std::vector<torch::Tensor> volumetric_rendering_backward(
-    torch::Tensor accumulated_weight, 
-    torch::Tensor accumulated_depth, 
-    torch::Tensor accumulated_color, 
-    torch::Tensor grad_weight, 
-    torch::Tensor grad_depth, 
-    torch::Tensor grad_color, 
-    torch::Tensor packed_info, 
-    torch::Tensor starts, 
-    torch::Tensor ends, 
-    torch::Tensor sigmas, 
-    torch::Tensor rgbs,
-    torch::Tensor opacities
-) {
-    DEVICE_GUARD(packed_info);
-    const uint32_t n_rays = packed_info.size(0);
-
-    const int threads = 256;
-    const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads);
-
-    // outputs
-    torch::Tensor grad_sigmas = torch::zeros(sigmas.sizes(), sigmas.options()); 
-    torch::Tensor grad_rgbs = torch::zeros(rgbs.sizes(), rgbs.options()); 
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        sigmas.scalar_type(),
-        "volumetric_rendering_backward",
-        ([&]
-         { volumetric_rendering_backward_kernel<scalar_t><<<blocks, threads>>>(
-                n_rays,
-                packed_info.data_ptr<int>(), 
-                starts.data_ptr<scalar_t>(),
-                ends.data_ptr<scalar_t>(),
-                sigmas.data_ptr<scalar_t>(),
-                rgbs.data_ptr<scalar_t>(),
-                opacities.data_ptr<scalar_t>(),
-                accumulated_weight.data_ptr<scalar_t>(),
-                accumulated_depth.data_ptr<scalar_t>(),
-                accumulated_color.data_ptr<scalar_t>(),
-                grad_weight.data_ptr<scalar_t>(),
-                grad_depth.data_ptr<scalar_t>(),
-                grad_color.data_ptr<scalar_t>(),
-                grad_sigmas.data_ptr<scalar_t>(),
-                grad_rgbs.data_ptr<scalar_t>()
-            ); 
-        }));
-
-    return {grad_sigmas, grad_rgbs};
-}
\ No newline at end of file
diff --git a/nerfactory/fields/density_fields/__init__.py b/nerfactory/fields/density_fields/__init__.py
deleted file mode 100644
index 02655775e..000000000
--- a/nerfactory/fields/density_fields/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright 2022 The Plenoptix Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
diff --git a/nerfactory/fields/density_fields/base.py b/nerfactory/fields/density_fields/base.py
deleted file mode 100644
index 02655775e..000000000
--- a/nerfactory/fields/density_fields/base.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright 2022 The Plenoptix Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
diff --git a/nerfactory/fields/density_fields/density_grid.py b/nerfactory/fields/density_fields/density_grid.py
deleted file mode 100644
index 6f8900b7c..000000000
--- a/nerfactory/fields/density_fields/density_grid.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# Copyright 2022 The Plenoptix Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Code to implement the density grid.
-"""
-
-from __future__ import annotations
-
-from typing import Callable, List, Tuple, Union
-
-import torch
-from torch import nn
-from torchtyping import TensorType
-
-import nerfactory.cuda as nerfactory_cuda
-from nerfactory.cameras.rays import RaySamples
-
-
-def _create_grid_coords(resolution: int, device: Union[torch.device, str] = "cpu") -> TensorType["n_coords", 3]:
-    """Create 3D grid coordinates. Returns all grid coordinates with shape [res * res * res, 3].
-
-    Args:
-        resolution: The 3D resolution of the grid.
-        device: Which device you want the returned tensor to live on.
-    """
-    arrange = torch.arange(resolution, device=device)
-    grid_x, grid_y, grid_z = torch.meshgrid([arrange, arrange, arrange], indexing="ij")
-    coords = torch.stack([grid_x, grid_y, grid_z])  # [3, res, res, res]
-    coords = coords.reshape(3, -1).t().contiguous()  # [N, 3]
-    return coords
-
-
-class DensityGrid(nn.Module):
-    """Cascaded multi-res density grids.
-
-    The DensityGrid contains multi-level (`num_cascades`) grids from finest (0) to the coarsest (num_cascades-1).
-    All levels of grid have the same `resolution` but cover different regions in the space. The finest grid is the
-    smallest one, which has scale of `base_scale` and centered at `center`. Every next level of grid grows by 2x
-    on the scale but all shares the same center.
-
-    For example:
-    - `center=0`, `base_scale=3`, `num_cascades=1` would create one grid (lvl=0) covers
-            the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]].
-    - `center=0`, `base_scale=3`, `num_cascades=2` would create one grid (lvl=0) covers
-            the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]],
-    and one more grid (lvl=1) covers the aabb [[-3, -3, -3], [3, 3, 3]].
-
-    Normally you can simply set the `center` to the scene center, and `base_scale` to be the scene scale, and use
-    `num_cascades=1` to create a single level density grid that covers the scene.
-
-    If, for example, you cares more about the central region of the scene and want to have more fine-grained density
-    grid at central region, you can set the `base_scale` to be the scale of the central region that you care. And adjust
-    `num_cascades` so that the coarsest level grid covers your entire scene. You should follow this rule so that the
-    multi-level grids can cover your scene:
-
-    scene_scale <= base_scale * 2 ** (num_cascades - 1)
-
-    TODO(ruilongli): support `center` and `base_scale` with 3-dim.
-
-    Args:
-        center: Center of all the grids in the world space. Defaults to 0.
-        base_scale: Center of the scale of the base level grid. Defaults to 1.
-        num_cascades: Number of the cascaded multi-res levels. Defaults to 1.
-        resolution: Resolution of the grid. Defaults to 128.
-        update_every_num_iters: How frequently to update the grid values. Defaults to 16.
-    """
-
-    density_grid: TensorType["num_cascades", "resolution_cubed"]
-    mean_density: TensorType[1]
-
-    def __init__(
-        self, center: float, base_scale: float, num_cascades: int, resolution: int, update_every_num_iters: int
-    ) -> None:
-        super().__init__()
-        self.center = center
-        self.base_scale = base_scale
-        self.num_cascades = num_cascades  # the number of levels (i.e, cascades)
-        self.resolution = resolution
-        self.update_every_num_iters = update_every_num_iters
-
-        density_grid = torch.zeros([self.num_cascades] + [self.resolution**3])
-        self.register_buffer("density_grid", density_grid)
-
-        density_bitfield = torch.zeros([self.num_cascades * self.resolution**3 // 8], dtype=torch.uint8)
-        self.register_buffer("density_bitfield", density_bitfield)
-
-        # Integer grid coords / indices that do not related to cascades
-        grid_coords = _create_grid_coords(self.resolution)
-        # TODO(ruilongli): hacky way for now. support cpu version of morton3D?
-        grid_indices = nerfactory_cuda.morton3D(grid_coords.to("cuda:0")).to("cpu")
-        self.register_buffer("grid_coords", grid_coords)
-        self.register_buffer("grid_indices", grid_indices)
-
-        mean_density = torch.zeros(())
-        self.register_buffer("mean_density", mean_density)
-
-    @torch.no_grad()
-    def get_all_cells(self) -> List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]:
-        """Returns a list of all grid indices and grid coordinates of all cascade levels."""
-        return [(self.grid_indices, self.grid_coords)] * self.num_cascades
-
-    @torch.no_grad()
-    def sample_uniform_and_occupied_cells(self, n: int) -> List[Tuple[TensorType["n"], TensorType["n", 3]]]:
-        """Samples both n uniform and occupied cells (per cascade).
-        Returns a tuple of (sampled grid indices, grid coordinates of all cascade levels).
-
-        Args:
-            n: Number of cells to be sampled.
-        """
-        device = self.density_grid.device
-
-        cells = []
-        for c in range(self.num_cascades):
-            uniform_coords = torch.randint(self.resolution, (n, 3), device=device)
-            uniform_indices = nerfactory_cuda.morton3D(uniform_coords)
-
-            occupied_indices = torch.nonzero(self.density_grid[c] > 0)[:, 0]
-            if n < len(occupied_indices):
-                selector = torch.randint(len(occupied_indices), (n,), device=device)
-                occupied_indices = occupied_indices[selector]
-            occupied_coords = nerfactory_cuda.morton3D_invert(occupied_indices)
-
-            # concatenate
-            coords = torch.cat([uniform_coords, occupied_coords], dim=0)
-            indices = torch.cat([uniform_indices, occupied_indices], dim=0)
-            cells.append((indices, coords))
-        return cells
-
-    @torch.no_grad()
-    def update_density_grid(
-        self,
-        density_eval_func: Callable,
-        step: int,
-        density_threshold: float = 2,  # 0.01 / (SQRT3 / 1024 * 3)
-        decay: float = 0.95,
-    ) -> None:
-        """Update the density grid in EMA way.
-
-        Args:
-            density_eval_func: A Callable function that takes in sample positions (N, 3) and
-                returns densities (N, 1).
-            step: The current training step. Instant-NGP has a warmup stage where all cells are updated.
-                After certain amount of steps (256), it speeds up by only update sampled cells.
-            density_threshold: The threshold to prune cells. Recommended to calculate it using this rule:
-                `weight_threshold / dt`. For example for Instant-NGP on Lego, the minimum step size `dt` is
-                `3` (scale of the scene) * `sqrt(3)` / `1024` (number of samples). And if we want to the weight
-                threshold to be `0.01`, we will get `0.01 / (3 * sqrt(3) / 1024) ~= 2.` for density threshold.
-            decay: EMA decay for density updating.
-        """
-        # create temporary grid
-        tmp_grid = -torch.ones_like(self.density_grid)
-        if step < 256:
-            cells = self.get_all_cells()
-        else:
-            N = self.resolution**3 // 4
-            cells = self.sample_uniform_and_occupied_cells(N)
-
-        # infer sigmas
-        for mip_level in range(self.num_cascades):
-            mip_scale = 2 ** (-mip_level) / self.base_scale
-            indices, coords = cells[mip_level]
-            # `coords` denotes the i-th cell. It's a poor naming here.
-            # the actually coordinates in world space x has mapping to
-            # the cell like this:
-            # x \in [0, 1/res] maps to 0-th cell.
-            # x \in [1 - 1/res, 1] maps to (res-1)-th cell.
-            x = (coords + torch.rand_like(coords.float())) / self.resolution
-            x = (x - 0.5) / mip_scale + self.center
-            tmp_grid[mip_level, indices] = density_eval_func(x).squeeze(-1)
-
-        # ema update
-        valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0)
-        self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask])
-        self.mean_density.data = self.density_grid.mean()
-
-        # pack to bitfield
-        self.density_bitfield.data = nerfactory_cuda.packbits(
-            self.density_grid, min(self.mean_density.item(), density_threshold)  # type: ignore
-        )
-
-        # TODO(ruilongli): max pooling? https://github.com/NVlabs/instant-ngp/blob/master/src/testbed_nerf.cu#L578
-
-    @torch.no_grad()
-    def get_densities(self, ray_samples: RaySamples) -> TensorType[..., 1]:
-        """Get the density values stored in the density grid.
-
-        Args:
-            ray_samples: ray samples to get density for
-        """
-        positions = ray_samples.frustums.get_positions()
-        deltas = ray_samples.frustums.ends - ray_samples.frustums.starts
-
-        _mip_levels, indices, _occupancies = nerfactory_cuda.occupancy_query(
-            positions.reshape(-1, 3),
-            deltas.reshape(-1),
-            self.center,
-            self.base_scale,
-            self.num_cascades,
-            self.resolution,
-            self.density_bitfield,
-        )
-
-        densities = self.density_grid.reshape(-1)[indices.long()]
-        densities = densities.view(*positions.shape[:-1], 1)
-        return densities
-
-    def forward(self, x):  # pylint: disable=no-self-use
-        """Not implemented."""
-        raise RuntimeError("Shouldn't be called!")
diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py
index 97a80dd79..81bfc130a 100644
--- a/nerfactory/models/base.py
+++ b/nerfactory/models/base.py
@@ -31,7 +31,6 @@ from nerfactory.cameras.rays import RayBundle
 from nerfactory.configs.base import InstantiateConfig
 from nerfactory.configs.utils import to_immutable_dict
 from nerfactory.datamanagers.structs import SceneBounds
-from nerfactory.fields.density_fields.density_grid import DensityGrid
 from nerfactory.models.modules.scene_colliders import NearFarCollider
 from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes
 
@@ -48,18 +47,6 @@ class ModelConfig(InstantiateConfig):
     collider_params: Optional[Dict[str, float]] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0})
     """parameters to instantiate scene collider with"""
     loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0})
-    """Loss specific weights."""
-    enable_density_field: bool = False
-    """Whether to create a density field to filter samples."""
-    density_field_params: Dict[str, Any] = to_immutable_dict(
-        {
-            "center": 0.0,  # simply set it as the center of the scene bbox
-            "base_scale": 3.0,  # simply set it as the scale of the scene bbox
-            "num_cascades": 1,  # if using more than 1 cascade, the `base_scale` can be smaller than scene scale.
-            "resolution": 128,
-            "update_every_num_iters": 16,
-        }
-    )
     """parameters to instantiate density field with"""
     eval_num_rays_per_chunk: int = 4096
     """specifies number of rays per chunk during eval"""
@@ -88,7 +75,6 @@ class Model(nn.Module):
         self.config = config
         self.scene_bounds = scene_bounds
         self.num_train_data = num_train_data
-        self.density_field = None
         self.kwargs = kwargs
         self.collider = None
         self.populate_modules()  # populate the modules
@@ -112,14 +98,6 @@ class Model(nn.Module):
         # default instantiates optional modules that are common among many networks
         # NOTE: call `super().populate_modules()` in subclasses
 
-        if self.config.enable_density_field:
-            self.density_field = DensityGrid(
-                center=self.config.density_field_params["center"],
-                base_scale=self.config.density_field_params["base_scale"],
-                num_cascades=self.config.density_field_params["num_cascades"],
-                resolution=self.config.density_field_params["resolution"],
-                update_every_num_iters=self.config.density_field_params["update_every_num_iters"],
-            )
         if self.config.enable_collider:
             self.collider = NearFarCollider(
                 near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"]
diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py
index c27e4d891..0f08b5529 100644
--- a/nerfactory/models/modules/ray_sampler.py
+++ b/nerfactory/models/modules/ray_sampler.py
@@ -26,7 +26,6 @@ from torch import nn
 from torchtyping import TensorType
 
 from nerfactory.cameras.rays import Frustums, RayBundle, RaySamples
-from nerfactory.fields.density_fields.density_grid import DensityGrid
 
 
 class Sampler(nn.Module):
@@ -34,34 +33,22 @@ class Sampler(nn.Module):
 
     Args:
         num_samples: number of samples to take
-        density_field: density grid specifying weighting of
     """
 
     def __init__(
         self,
         num_samples: Optional[int] = None,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 0.01,
     ) -> None:
         super().__init__()
         self.num_samples = num_samples
-        self.density_field = density_field
-        self.weight_threshold = weight_threshold
 
     @abstractmethod
     def generate_ray_samples(self) -> RaySamples:
         """Generate Ray Samples"""
 
     def forward(self, *args, **kwargs) -> RaySamples:
-        """Generate ray samples with optional density filtering"""
-        ray_samples = self.generate_ray_samples(*args, **kwargs)
-        if self.density_field is not None:
-            densities = self.density_field.get_densities(ray_samples)
-            deltas = torch.clamp(ray_samples.frustums.ends - ray_samples.frustums.starts, min=1e-10)
-            density_threshold = torch.clamp(self.weight_threshold / deltas, max=self.density_field.mean_density)
-            valid_mask = densities > density_threshold
-            ray_samples.set_valid_mask(valid_mask & ray_samples.valid_mask)
-        return ray_samples
+        """Generate ray samples"""
+        return self.generate_ray_samples(*args, **kwargs)
 
 
 class SpacedSampler(Sampler):
@@ -72,8 +59,6 @@ class SpacedSampler(Sampler):
         spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).
         spacing_fn_inv: The inverse of spacing_fn.
         train_stratified: Use stratified sampling during training. Defults to True
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
     """
 
     def __init__(
@@ -83,10 +68,8 @@ class SpacedSampler(Sampler):
         num_samples: Optional[int] = None,
         train_stratified=True,
         single_jitter=False,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
-        super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
+        super().__init__(num_samples=num_samples)
         self.train_stratified = train_stratified
         self.single_jitter = single_jitter
         self.spacing_fn = spacing_fn
@@ -149,8 +132,6 @@ class UniformSampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
     """
 
     def __init__(
@@ -158,8 +139,6 @@ class UniformSampler(SpacedSampler):
         num_samples: Optional[int] = None,
         train_stratified=True,
         single_jitter=False,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
         super().__init__(
             num_samples=num_samples,
@@ -167,8 +146,6 @@ class UniformSampler(SpacedSampler):
             spacing_fn_inv=lambda x: x,
             train_stratified=train_stratified,
             single_jitter=single_jitter,
-            density_field=density_field,
-            weight_threshold=weight_threshold,
         )
 
 
@@ -178,8 +155,6 @@ class LinearDisparitySampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
     """
 
     def __init__(
@@ -187,8 +162,6 @@ class LinearDisparitySampler(SpacedSampler):
         num_samples: Optional[int] = None,
         train_stratified=True,
         single_jitter=False,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
         super().__init__(
             num_samples=num_samples,
@@ -196,8 +169,6 @@ class LinearDisparitySampler(SpacedSampler):
             spacing_fn_inv=lambda x: 1 / x,
             train_stratified=train_stratified,
             single_jitter=single_jitter,
-            density_field=density_field,
-            weight_threshold=weight_threshold,
         )
 
 
@@ -207,8 +178,6 @@ class SqrtSampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
     """
 
     def __init__(
@@ -216,8 +185,6 @@ class SqrtSampler(SpacedSampler):
         num_samples: Optional[int] = None,
         train_stratified=True,
         single_jitter=False,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
         super().__init__(
             num_samples=num_samples,
@@ -225,8 +192,6 @@ class SqrtSampler(SpacedSampler):
             spacing_fn_inv=lambda x: x**2,
             train_stratified=train_stratified,
             single_jitter=single_jitter,
-            density_field=density_field,
-            weight_threshold=weight_threshold,
         )
 
 
@@ -236,8 +201,6 @@ class LogSampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
     """
 
     def __init__(
@@ -245,8 +208,6 @@ class LogSampler(SpacedSampler):
         num_samples: Optional[int] = None,
         train_stratified=True,
         single_jitter=False,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
         super().__init__(
             num_samples=num_samples,
@@ -254,8 +215,6 @@ class LogSampler(SpacedSampler):
             spacing_fn_inv=torch.exp,
             train_stratified=train_stratified,
             single_jitter=single_jitter,
-            density_field=density_field,
-            weight_threshold=weight_threshold,
         )
 
 
@@ -267,8 +226,6 @@ class UniformLinDispPiecewiseSampler(SpacedSampler):
     Args:
         num_samples: Number of samples per ray
         train_stratified: Use stratified sampling during training. Defults to True
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
     """
 
     def __init__(
@@ -276,8 +233,6 @@ class UniformLinDispPiecewiseSampler(SpacedSampler):
         num_samples: Optional[int] = None,
         train_stratified=True,
         single_jitter=False,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
         super().__init__(
             num_samples=num_samples,
@@ -285,8 +240,6 @@ class UniformLinDispPiecewiseSampler(SpacedSampler):
             spacing_fn_inv=lambda x: torch.where(x < 0.5, 2 * x, 1 / (2 - 2 * x)),
             train_stratified=train_stratified,
             single_jitter=single_jitter,
-            density_field=density_field,
-            weight_threshold=weight_threshold,
         )
 
 
@@ -297,8 +250,6 @@ class PDFSampler(Sampler):
         num_samples: Number of samples per ray
         train_stratified: Randomize location within each bin during training.
         include_original: Add original samples to ray.
-        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
-        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
         histogram_padding: Amount to weights prior to computing PDF.
     """
 
@@ -307,11 +258,9 @@ class PDFSampler(Sampler):
         num_samples: Optional[int] = None,
         train_stratified: bool = True,
         include_original: bool = True,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
         histogram_padding: float = 0.01,
     ) -> None:
-        super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
+        super().__init__(num_samples=num_samples)
         self.train_stratified = train_stratified
         self.include_original = include_original
         self.histogram_padding = histogram_padding
@@ -490,9 +439,6 @@ class VolumetricSampler(Sampler):
             The ray_indices contains the indices of the rays that each sample belongs to.
         """
 
-        if self.density_field is not None:
-            raise ValueError("NGPSpacedSampler does not support a density field")
-
         rays_o = ray_bundle.origins.contiguous()
         rays_d = ray_bundle.directions.contiguous()
         if ray_bundle.camera_indices is not None:
@@ -550,10 +496,8 @@ class ProposalNetworkSampler(Sampler):
         num_nerf_samples_per_ray: int = 32,
         num_proposal_network_iterations: int = 2,
         single_jitter: bool = True,
-        density_field: Optional[DensityGrid] = None,
-        weight_threshold: float = 1e-2,
     ) -> None:
-        super().__init__(density_field=density_field, weight_threshold=weight_threshold)
+        super().__init__()
         self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
         self.num_nerf_samples_per_ray = num_nerf_samples_per_ray
         self.num_proposal_network_iterations = num_proposal_network_iterations
diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py
index fccfdcb40..7737671ac 100644
--- a/nerfactory/models/tensorf.py
+++ b/nerfactory/models/tensorf.py
@@ -151,10 +151,8 @@ class TensoRFModel(Model):
         )
 
         # samplers
-        self.sampler_uniform = UniformSampler(
-            num_samples=self.config.num_coarse_samples, density_field=self.density_field
-        )
-        self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, density_field=self.density_field)
+        self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)
+        self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples)
 
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py
index 85ffbdfd6..15f378f56 100644
--- a/nerfactory/models/vanilla_nerf.py
+++ b/nerfactory/models/vanilla_nerf.py
@@ -39,11 +39,6 @@ from nerfactory.renderers.renderers import (
     RGBRenderer,
 )
 from nerfactory.utils import colors, misc, visualization
-from nerfactory.utils.callbacks import (
-    TrainingCallback,
-    TrainingCallbackAttributes,
-    TrainingCallbackLocation,
-)
 
 
 class NeRFModel(Model):
@@ -66,23 +61,6 @@ class NeRFModel(Model):
             **kwargs,
         )
 
-    def get_training_callbacks(
-        self, training_callback_attributes: TrainingCallbackAttributes  # pylint: disable=unused-argument
-    ) -> List[TrainingCallback]:
-        if self.field_coarse is None:
-            raise ValueError("populate fields must be called before get_training_callbacks.")
-        callbacks = []
-        if self.density_field is not None:
-            callbacks = [
-                TrainingCallback(
-                    where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                    update_every_num_iters=self.density_field.update_every_num_iters,
-                    func=self.density_field.update_density_grid,
-                    kwargs={"density_eval_func": self.field_coarse.density_fn},  # type: ignore
-                )
-            ]
-        return callbacks
-
     def populate_modules(self):
         """Set the fields and modules"""
         super().populate_modules()
@@ -99,10 +77,8 @@ class NeRFModel(Model):
         self.field_fine = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding)
 
         # samplers
-        self.sampler_uniform = UniformSampler(
-            num_samples=self.config.num_coarse_samples, density_field=self.density_field
-        )
-        self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, density_field=self.density_field)
+        self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)
+        self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples)
 
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
diff --git a/pyproject.toml b/pyproject.toml
index 7294e4363..585fb2571 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -97,7 +97,7 @@ line-length = 120
 max-line-length = 120
 generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"]
 good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$"
-ignore-paths = ["scripts/colmap2nerf.py", "^nerfactory/cuda/.*$", "^tests/cuda/.*$"]
+ignore-paths = ["scripts/colmap2nerf.py"]
 jobs = 0
 
 disable = [
@@ -113,7 +113,7 @@ disable = [
 
 #pytest
 [tool.pytest.ini_options]
-addopts = "-n=4 --typeguard-packages=nerfactory --torchtyping-patch-typeguard --disable-warnings --ignore=tests/cuda"
+addopts = "-n=4 --typeguard-packages=nerfactory --torchtyping-patch-typeguard --disable-warnings"
 testpaths = [
     "tests",
 ]
@@ -124,7 +124,7 @@ include = ["nerfactory"]
 exclude = ["**/node_modules",
     "**/__pycache__",
 ]
-ignore = ["nerfactory/viewer", "nerfactory/cuda"]
+ignore = ["nerfactory/viewer"]
 defineConstant = { DEBUG = true }
 
 reportMissingImports = true
diff --git a/tests/cuda/test.py b/tests/cuda/test.py
deleted file mode 100644
index ec34acb52..000000000
--- a/tests/cuda/test.py
+++ /dev/null
@@ -1,3 +0,0 @@
-if __name__ == "__main__":
-    # TODO(ruilongli): write tests
-    pass
diff --git a/tests/graphs/modules/test_density_grid.py b/tests/graphs/modules/test_density_grid.py
deleted file mode 100644
index c99329cc5..000000000
--- a/tests/graphs/modules/test_density_grid.py
+++ /dev/null
@@ -1,26 +0,0 @@
-"""
-Code to test the density grid.
-"""
-
-# import torch
-
-# from nerfactory.fields.density_fields.density_grid import DensityGrid
-
-
-def test_density_grid():
-    """For testing the density grid."""
-    # TODO(ruilongli): currently the DensityGrid only supports GPU init.
-    # after we support for CPU init then we re-enable this test
-
-    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    # num_cascades = 1
-    # resolution = 128
-    # aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32).to(device)
-    # grid = DensityGrid(num_cascades, resolution, aabb)
-    # grid.to(device)
-    # density_eval_func = lambda x: x[..., 0:1]  # xyz (..., 3) => density (...,)
-    # grid.update_density_grid(density_eval_func=density_eval_func, step=0)
-
-
-if __name__ == "__main__":
-    test_density_grid()