class_docstr: NeRF solver class.
  
  Args:
      device: device for class tensors: Union[str, Device]
      dtype: type for all floating point calculations: torch.dtype
class_docstr_tok: ["NeRF","solver","class",".","Args",":","device",":","device","for","class","tensors",":","Union","[","str",",","Device","]","dtype",":","type","for","all","floating","point","calculations",":","torch.dtype"]
docstr: Trains the NeRF model one epoch.
  
  1) A dataset of rays is initialized, and sent over to a data loader.
  2) The data loader sample a batch of rays randomly, and runs them through the NeRF model,
  to predict ray associated rgb model values.
  3) The model rgb is compared with the image pixel rgb, and the loss between the two is back
  propagated to update the model weights.
  
  Implemented steps:
      - Create an object of class RayDataset
      - Initialize ray dataset with group of images on disk, and number of rays to randomly sample
      - Initialize a data loader with batch size info
      - Iterate over data loader
      -- Reset optimizer
      -- Run ray batch through Nerf model
      -- Find loss
      -- Back propagate loss
      -- Optimizer step
  
  Returns:
      Average psnr over all epoch rays.
docstr_tok: ["Trains","the","NeRF","model","one","epoch",".","1",")","A","dataset","of","rays","is","initialized",",","and","sent","over","to","a","data","loader",".","2",")","The","data","loader","sample","a","batch","of","rays","randomly",",","and","runs","them","through","the","NeRF","model",",","to","predict","ray","associated","rgb","model","values",".","3",")","The","model","rgb","is","compared","with","the","image","pixel","rgb",",","and","the","loss","between","the","two","is","back","propagated","to","update","the","model","weights",".","Implemented","steps",":","-","Create","an","object","of","class","RayDataset","-","Initialize","ray","dataset","with","group","of","images","on","disk",",","and","number","of","rays","to","randomly","sample","-","Initialize","a","data","loader","with","batch","size","info","-","Iterate","over","data","loader","--","Reset","optimizer","--","Run","ray","batch","through","Nerf","model","--","Find","loss","--","Back","propagate","loss","--","Optimizer","step","Returns",":","Average","psnr","over","all","epoch","rays","."]
code: def _train_one_epoch(self) -> float:
      r"""Trains the NeRF model one epoch.
  
      1) A dataset of rays is initialized, and sent over to a data loader.
      2) The data loader sample a batch of rays randomly, and runs them through the NeRF model,
      to predict ray associated rgb model values.
      3) The model rgb is compared with the image pixel rgb, and the loss between the two is back
      propagated to update the model weights.
  
      Implemented steps:
          - Create an object of class RayDataset
          - Initialize ray dataset with group of images on disk, and number of rays to randomly sample
          - Initialize a data loader with batch size info
          - Iterate over data loader
          -- Reset optimizer
          -- Run ray batch through Nerf model
          -- Find loss
          -- Back propagate loss
          -- Optimizer step
  
      Returns:
          Average psnr over all epoch rays.
      """
      KORNIA_CHECK(
          self._nerf_model is not None,
          "The model should be a NeRF model.",
      )
      KORNIA_CHECK(
          self._nerf_optimizer is not None,
          "The optimizer should be an Adam optimizer.",
      )
      KORNIA_CHECK(
          self._cameras is not None,
          "The camera should be a PinholeCamera.",
      )
      KORNIA_CHECK(
          self._imgs is not None,
          "The images should be a list of tensors.",
      )
      KORNIA_CHECK(
          self._num_img_rays is not None,
          "The number of images of Ray should be a tensor.",
      )
  
      # TODO: refactor and so that the constructor receives the correct types
      cameras: PinholeCamera = cast(PinholeCamera, self._cameras)
      num_img_rays: Tensor = cast(Tensor, self._num_img_rays)
      images = cast(Images, self._imgs)
      nerf_model: NerfModel = cast(NerfModel, self._nerf_model)
      nerf_optimizer: optim.Optimizer = cast(
          optim.Optimizer, self._nerf_optimizer
      )
  
      # create the dataset and data loader
      ray_dataset = RayDataset(
          cameras,
          self._min_depth,
          self._max_depth,
          self._ndc,
          device=self._device,
          dtype=self._dtype,
      )
      ray_dataset.init_ray_dataset(num_img_rays)
      ray_dataset.init_images_for_training(
          images
      )  # FIXME: Do we need to load the same images on each Epoch?
  
      # data loader
      ray_data_loader = instantiate_ray_dataloader(
          ray_dataset, self._batch_size, shuffle=True
      )
  
      total_psnr: Tensor = torch.tensor(
          0.0, device=self._device, dtype=self._dtype
      )
  
      i_batch: float = 0
  
      for origins, directions, rgbs in ray_data_loader:
          rgbs_model = nerf_model(origins, directions)
          loss = F.mse_loss(rgbs_model, rgbs)
  
          total_psnr += psnr(rgbs_model, rgbs, 1.0)
  
          nerf_optimizer.zero_grad()
          loss.backward()
          nerf_optimizer.step()
  
          i_batch += 1
  
      return float(total_psnr / (i_batch + 1))
  
imports: from __future__ import annotations
  import logging
  from datetime import datetime
  from typing import cast
  import torch
  import torch.nn.functional
  from torch import optim
  from kornia.core import Module, Tensor, tensor
  from kornia.core.check import KORNIA_CHECK
  from kornia.geometry.camera import PinholeCamera
  from kornia.metrics import psnr
  from kornia.nerf.core import Images
  from kornia.nerf.data_utils import RayDataset, instantiate_ray_dataloader
  from kornia.nerf.nerf_model import NerfModel
  from kornia.utils import deprecated
label_desc: Use image node_id 5 for calling the NerfSolver obj's underlying member method code with example usage: obj._train_one_epoch() and returns: float