class_docstr: Original NeRF network.
class_docstr_tok: ["Original","NeRF","network","."]
docstr: Training function for coarse and fine networks.
  
  Args:
    r_org: A tensor of shape `[B, N, 3]` where B is the batch size, N is the
      number of rays.
    r_dir: A tensor of shape `[B, N, 3]` where B is the batch size, N is the
      number of rays.
    gt_rgb: A tensor of shape `[B, N, 3]` where B is the batch size, N is the
      number of rays.
  
  Returns:
    A scalar.
docstr_tok: ["Training","function","for","coarse","and","fine","networks",".","Args",":","r_org",":","A","tensor","of","shape","`","[","B",",","N",",","3","]","`","where","B","is","the","batch","size",",","N","is","the","number","of","rays",".","r_dir",":","A","tensor","of","shape","`","[","B",",","N",",","3","]","`","where","B","is","the","batch","size",",","N","is","the","number","of","rays",".","gt_rgb",":","A","tensor","of","shape","`","[","B",",","N",",","3","]","`","where","B","is","the","batch","size",",","N","is","the","number","of","rays",".","Returns",":","A","scalar","."]
code: def train_step(self, r_org, r_dir, gt_rgb):
      """Training function for coarse and fine networks.
  
      Args:
        r_org: A tensor of shape `[B, N, 3]` where B is the batch size, N is the
          number of rays.
        r_dir: A tensor of shape `[B, N, 3]` where B is the batch size, N is the
          number of rays.
        gt_rgb: A tensor of shape `[B, N, 3]` where B is the batch size, N is the
          number of rays.
  
      Returns:
        A scalar.
      """
      with tf.GradientTape() as tape:
          rgb_fine, rgb_coarse, _, _ = self.inference(r_org, r_dir)
          rgb_coarse_loss = utils.l2_loss(rgb_coarse, gt_rgb)
          rgb_fine_loss = utils.l2_loss(rgb_fine, gt_rgb)
          total_loss = rgb_coarse_loss + rgb_fine_loss
      gradients = tape.gradient(total_loss, self.network_vars)
      self.optimizer_network.apply_gradients(
          zip(gradients, self.network_vars)
      )
  
      with self.summary_writer.as_default():
          step = self.global_step
          tf.summary.scalar("total_loss", total_loss, step=step)
          tf.summary.scalar("rgb_loss_f", rgb_fine_loss, step=step)
          tf.summary.scalar("rgb_loss_c", rgb_coarse_loss, step=step)
      return total_loss
  
file_path: graphics/tensorflow_graphics/projects/radiance_fields/nerf/model.py
imports: from absl import logging
  import numpy
  import tensorflow
  import tensorflow_graphics.geometry.representation.ray
  import tensorflow_graphics.math.feature_representation
  import tensorflow_graphics.projects.radiance_fields.nerf.layers
  import tensorflow_graphics.projects.radiance_fields.utils
  import tensorflow_graphics.rendering.volumetric.ray_radiance
label_desc: Use image node_id 10 for calling the NeRF obj's underlying member method code with example usage: obj.train_step(r_org, r_dir, gt_rgb) and returns: total_loss