diff --git a/README.md b/README.md
index 35dd0886..50d0960f 100644
--- a/README.md
+++ b/README.md
@@ -62,10 +62,13 @@ Performance on test set:
 | Our train time & test FPS | 43min; 0.15FPS | 41min; 0.4FPS |
 
 
+Note the numbers here are tested with version `v0.0.8`
+
+<!-- 
 ## Tips:
 
 1. sample rays over all images per iteration (`batch_over_images=True`) is better: `PSNR 33.31 -> 33.75`.
 2. make use of scheduler (`MultiStepLR(optimizer, milestones=[20000, 30000], gamma=0.1)`) to adjust learning rate gives: `PSNR 33.75 -> 34.40`.
 3. increasing chunk size (`chunk: 8192 -> 81920`) during inference gives speedup: `FPS 4.x -> 6.2`
 4. random bkgd color (`color_bkgd_aug="random"`) for the `Lego` scene actually hurts: `PNSR 35.42 -> 34.38`
-
+ -->
diff --git a/nerfacc/cuda/csrc/volumetric_marching.cu b/nerfacc/cuda/csrc/volumetric_marching.cu
index a3b8ebfc..4c9b2c15 100644
--- a/nerfacc/cuda/csrc/volumetric_marching.cu
+++ b/nerfacc/cuda/csrc/volumetric_marching.cu
@@ -7,7 +7,6 @@ inline __device__ int cascaded_grid_idx_at(
     const int resx, const int resy, const int resz, 
     const float* aabb
 ) {
-    // TODO(ruilongli): if the x, y, z is outside the aabb, it will be clipped into aabb!!! We should just return false
     int ix = (int)(((x - aabb[0]) / (aabb[3] - aabb[0])) * resx);
     int iy = (int)(((y - aabb[1]) / (aabb[4] - aabb[1])) * resy);
     int iz = (int)(((z - aabb[2]) / (aabb[5] - aabb[2])) * resz);
@@ -15,7 +14,6 @@ inline __device__ int cascaded_grid_idx_at(
     iy = __clamp(iy, 0, resy-1);
     iz = __clamp(iz, 0, resz-1);
     int idx = ix * resy * resz + iy * resz + iz;
-    // printf("(ix, iy, iz) = (%d, %d, %d)\n", ix, iy, iz);
     return idx;
 }
 
@@ -35,13 +33,16 @@ 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,
-    const int resx, const int resy, const int resz
+    const int resx, const int resy, const int resz,
+    const float* aabb
 ) { // dda like step
-    // TODO: warning: expression has no effect?
-    x, y, z = resx * x, resy * y, resz * z;
-    float tx = ((floorf(x + 0.5f + 0.5f * __sign(dir_x)) - x) * idir_x) / resx;
-    float ty = ((floorf(y + 0.5f + 0.5f * __sign(dir_y)) - y) * idir_y) / resy;
-    float tz = ((floorf(z + 0.5f + 0.5f * __sign(dir_z)) - z) * idir_z) / resz;
+    // TODO: this is ugly -- optimize this.
+    float _x = ((x - aabb[0]) / (aabb[3] - aabb[0])) * resx;
+    float _y = ((y - aabb[1]) / (aabb[4] - aabb[1])) * resy;
+    float _z = ((z - aabb[2]) / (aabb[5] - aabb[2])) * resz;
+    float tx = ((floorf(_x + 0.5f + 0.5f * __sign(dir_x)) - _x) * idir_x) / resx * (aabb[3] - aabb[0]);
+    float ty = ((floorf(_y + 0.5f + 0.5f * __sign(dir_y)) - _y) * idir_y) / resy * (aabb[4] - aabb[1]);
+    float tz = ((floorf(_z + 0.5f + 0.5f * __sign(dir_z)) - _z) * idir_z) / resz * (aabb[5] - aabb[2]);
     float t = min(min(tx, ty), tz);
     return fmaxf(t, 0.0f);
 }
@@ -51,11 +52,14 @@ inline __device__ float advance_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,
-    const int resx, const int resy, const int resz,
+    const int resx, const int resy, const int resz, const float* aabb,
     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, resx, resy, resz
+        x, y, z, 
+        dir_x, dir_y, dir_z, 
+        idir_x, idir_y, idir_z, 
+        resx, resy, resz, aabb
     );
     do {
         t += dt_min;
@@ -106,7 +110,6 @@ __global__ void marching_steps_kernel(
         const float x = ox + t_mid * dx;
         const float y = oy + t_mid * dy;
         const float z = oz + t_mid * dz;
-        
         if (grid_occupied_at(x, y, z, resx, resy, resz, aabb, occ_binary)) {
             ++j;
             // march to next sample
@@ -117,7 +120,7 @@ __global__ void marching_steps_kernel(
         else {
             // march to next sample
             t_mid = advance_to_next_voxel(
-                t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, resx, resy, resz, dt
+                t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, resx, resy, resz, aabb, dt
             );
             t0 = t_mid - dt * 0.5f;
             t1 = t_mid + dt * 0.5f;
@@ -192,7 +195,7 @@ __global__ void marching_forward_kernel(
         else {
             // march to next sample
             t_mid = advance_to_next_voxel(
-                t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, resx, resy, resz, dt
+                t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, resx, resy, resz, aabb, dt
             );
             t0 = t_mid - dt * 0.5f;
             t1 = t_mid + dt * 0.5f;
diff --git a/pyproject.toml b/pyproject.toml
index 0bb7e3e7..b16d4fd7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nerfacc"
-version = "0.0.8"
+version = "0.0.9"
 authors = [{name = "Ruilong", email = "ruilongli94@gmail.com"}]
 license = { text="MIT" }
 requires-python = ">=3.8"
diff --git a/tests/test_marching.py b/tests/test_marching.py
index 6e27e1c5..01fd661b 100644
--- a/tests/test_marching.py
+++ b/tests/test_marching.py
@@ -7,8 +7,9 @@
 
 
 def test_marching():
+    torch.manual_seed(42)
     scene_aabb = torch.tensor([0, 0, 0, 1, 1, 1], device=device).float()
-    scene_occ_binary = torch.ones((128 * 128 * 128), device=device).bool()
+    scene_occ_binary = torch.rand((128 * 128 * 128), device=device) > 0.5
     rays_o = torch.rand((10000, 3), device=device)
     rays_d = torch.randn((10000, 3), device=device)
     rays_d = rays_d / rays_d.norm(dim=-1, keepdim=True)