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))
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/nerf/nerf_solver.py