code: def update_step(
      self,
      epoch: int,
      global_step: int,
      decay: float = 0.95,
      S: int = 128,
      **kwargs,
  ):
      if global_step % self.update_interval != 0:
          return
  
      ### update density grid
      tmp_grid = -torch.ones_like(self.density_grid)
  
      X = torch.arange(
          self.grid_resolution,
          dtype=torch.int32,
          device=self.aabb.device,
      ).split(S)
      Y = torch.arange(
          self.grid_resolution,
          dtype=torch.int32,
          device=self.aabb.device,
      ).split(S)
      Z = torch.arange(
          self.grid_resolution,
          dtype=torch.int32,
          device=self.aabb.device,
      ).split(S)
  
      for xs in X:
          for ys in Y:
              for zs in Z:
                  # construct points
                  xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing="ij")
                  coords = torch.cat(
                      [
                          xx.reshape(-1, 1),
                          yy.reshape(-1, 1),
                          zz.reshape(-1, 1),
                      ],
                      dim=-1,
                  )  # [N, 3], in [0, 128)
                  indices = raymarching.morton3D(coords).long()  # [N]
                  xyzs = (
                      2 * coords.float() / (self.grid_resolution - 1)
                      - 1
                  )  # [N, 3] in [-1, 1]
  
                  # cascading
                  for cas in range(self.cascade):
                      bound = min(2**cas, self.bound)
                      half_grid_resolution = (
                          bound / self.grid_resolution
                      )
                      # scale to current cascade's resolution
                      cas_xyzs = xyzs * (bound - half_grid_resolution)
                      # add noise in [-hgs, hgs]
                      cas_xyzs += (
                          torch.rand_like(cas_xyzs) * 2 - 1
                      ) * half_grid_resolution
                      # query density
                      density = (
                          self.nerf.forward_density(cas_xyzs)
                          .reshape(-1)
                          .detach()
                      )
                      # assign
                      tmp_grid[cas, indices] = density
      # ema update
      valid_mask = self.density_grid >= 0
      self.density_grid[valid_mask] = torch.maximum(
          self.density_grid[valid_mask] * decay, tmp_grid[valid_mask]
      )
      self.mean_density = torch.mean(
          self.density_grid[valid_mask]
      ).item()
      self.iter_density += 1
  
      # convert to bitfield
      density_thresh = min(self.mean_density, self.density_thresh)
      self.density_bitfield = raymarching.packbits(
          self.density_grid, density_thresh, self.density_bitfield
      )
  
file_path: NeMo/nemo/collections/multimodal/modules/nerf/renderers/torchngp_volume_renderer.py
imports: import math
  import torch
  import torch.nn.functional
  import nemo.collections.multimodal.modules.nerf.utils.torch_ngp.raymarching
  from nemo.collections.multimodal.modules.nerf.materials.materials_base import ShadingEnum
  from nemo.collections.multimodal.modules.nerf.renderers.base_renderer import BaseRenderer
label_desc: Use image node_id 2 for calling the TorchNGPVolumeRenderer obj's underlying member method code with example usage: obj.update_step(epoch, global_step, decay, S) without return types