diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py
index 9ea9bdcc2..51f85e7aa 100644
--- a/nerfactory/cameras/rays.py
+++ b/nerfactory/cameras/rays.py
@@ -165,8 +165,6 @@ class RayBundle(TensorDataclass):
     """Rays Distance along ray to stop sampling"""
     valid_mask: Optional[TensorType[..., 1, bool]] = None
     """Rays that are valid"""
-    num_rays_per_chunk: int = 128
-    """Number of rays per chunk"""
     metadata: Optional[Dict[str, TensorType["num_rays", "latent_dims"]]] = None
     """Additional metadata or data needed for interpolation, will mimic shape of rays"""
 
diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py
index 0b0958067..842ea0976 100644
--- a/nerfactory/configs/base.py
+++ b/nerfactory/configs/base.py
@@ -182,8 +182,12 @@ class TrainerConfig(PrintableConfig):
 
     steps_per_save: int = 1000
     """number of steps between saves"""
-    steps_per_test: int = 500
-    """number of steps between eval"""
+    steps_per_eval_batch: int = 500
+    """number of steps between randomly sampled batches of rays"""
+    steps_per_eval_image: int = 500
+    """number of steps between single eval images"""
+    steps_per_eval_all_images: int = 25000
+    """number of steps between eval all images"""
     max_num_iterations: int = 1000000
     """maximum number of iterations to run"""
     mixed_precision: bool = False
@@ -223,9 +227,14 @@ class NerfactoryDataParserConfig(DataParserConfig):
     downscale_factor: int = 1
     """How much to downscale images. Defaults to 1."""
     scene_scale: float = 4.0
-    """How much to scale the scene. Defaults to 0.33"""
+    """How much to scale the scene. Defaults to 4.0"""
     orientation_method: Literal["pca", "up"] = "up"
     """The method to use for orientation. Either "pca" or "up"."""
+    train_split_percentage: float = 0.9
+    """
+    The percent of images to use for training.
+    The remaining images are for eval. Defaults to 0.9.
+    """
 
 
 @dataclass
@@ -351,8 +360,6 @@ class VanillaDataManagerConfig(InstantiateConfig):
     """number of images to sample during eval iteration"""
     eval_image_indices: Optional[Tuple[int, ...]] = (0,)
     """specifies the image indices to use during eval; if None, uses all"""
-    eval_num_rays_per_chunk: int = 4096
-    """specifies number of rays per chunk during eval"""
 
 
 @dataclass
@@ -394,6 +401,8 @@ class ModelConfig(InstantiateConfig):
         }
     )
     """parameters to instantiate density field with"""
+    eval_num_rays_per_chunk: int = 4096
+    """specifies number of rays per chunk during eval"""
 
 
 @dataclass
@@ -564,7 +573,9 @@ class Config(PrintableConfig):
             # also disable eval steps
             string = "Disabling eval iterations since viewer is enabled."
             CONSOLE.print(f"[bold red]{string}")
-            self.trainer.steps_per_test = self.trainer.max_num_iterations
+            self.trainer.steps_per_eval_batch = self.trainer.max_num_iterations
+            self.trainer.steps_per_eval_image = self.trainer.max_num_iterations
+            self.trainer.steps_per_eval_all_images = self.trainer.max_num_iterations
 
     def set_timestamp(self, timestamp: Optional[str] = None) -> None:
         """Make paths in our config more specific using a timestamp.
diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py
index 4734efbc8..f3e0de891 100644
--- a/nerfactory/configs/base_configs.py
+++ b/nerfactory/configs/base_configs.py
@@ -46,14 +46,10 @@ from nerfactory.models.vanilla_nerf import NeRFModel
 base_configs: Dict[str, Config] = {}
 base_configs["instant_ngp"] = Config(
     method_name="instant_ngp",
-    trainer=TrainerConfig(steps_per_test=500, steps_per_save=2000, mixed_precision=True),
+    trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True),
     pipeline=PipelineConfig(
-        datamanager=VanillaDataManagerConfig(
-            train_dataparser=BlenderDataParserConfig(),
-            train_num_rays_per_batch=8192,
-            eval_num_rays_per_chunk=8192,
-        ),
-        model=InstantNGPModelConfig(),
+        datamanager=VanillaDataManagerConfig(train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192),
+        model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192),
     ),
     optimizers={
         "fields": {
@@ -66,12 +62,10 @@ base_configs["instant_ngp"] = Config(
 base_configs["mipnerf_360"] = Config(
     experiment_name="mipnerf_360",
     method_name="mipnerf_360",
-    trainer=TrainerConfig(steps_per_test=200),
+    trainer=TrainerConfig(steps_per_eval_batch=200),
     pipeline=PipelineConfig(
         datamanager=VanillaDataManagerConfig(
-            train_dataparser=MipNerf360DataParserConfig(),
-            train_num_rays_per_batch=8192,
-            eval_num_rays_per_chunk=8192,
+            train_dataparser=MipNerf360DataParserConfig(), train_num_rays_per_batch=8192
         ),
         model=ModelConfig(
             _target=MipNerf360Model,
@@ -79,6 +73,7 @@ base_configs["mipnerf_360"] = Config(
             loss_coefficients={"ray_loss_coarse": 1.0, "ray_loss_fine": 1.0},
             num_coarse_samples=128,
             num_importance_samples=128,
+            eval_num_rays_per_chunk=8192,
         ),
     ),
 )
@@ -86,16 +81,13 @@ base_configs["mipnerf_360"] = Config(
 base_configs["mipnerf"] = Config(
     method_name="mipnerf",
     pipeline=PipelineConfig(
-        datamanager=VanillaDataManagerConfig(
-            train_dataparser=BlenderDataParserConfig(),
-            train_num_rays_per_batch=8192,
-            eval_num_rays_per_chunk=8192,
-        ),
+        datamanager=VanillaDataManagerConfig(train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192),
         model=ModelConfig(
             _target=MipNerfModel,
             loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0},
             num_coarse_samples=128,
             num_importance_samples=128,
+            eval_num_rays_per_chunk=8192,
         ),
     ),
 )
@@ -159,12 +151,8 @@ base_configs["compound"] = Config(
     method_name="compound",
     trainer=TrainerConfig(mixed_precision=True),
     pipeline=PipelineConfig(
-        datamanager=VanillaDataManagerConfig(
-            train_dataparser=BlenderDataParserConfig(),
-            train_num_rays_per_batch=8192,
-            eval_num_rays_per_chunk=8192,
-        ),
-        model=CompoundModelConfig(),
+        datamanager=VanillaDataManagerConfig(train_dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192),
+        model=CompoundModelConfig(eval_num_rays_per_chunk=8192),
     ),
     optimizers={
         "fields": {
diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py
index 95bd500d5..6c1722215 100644
--- a/nerfactory/datamanagers/base.py
+++ b/nerfactory/datamanagers/base.py
@@ -31,6 +31,7 @@ from nerfactory.configs import base as cfg
 from nerfactory.datamanagers.dataloaders import (
     CacheImageDataloader,
     FixedIndicesEvalDataloader,
+    RandIndicesEvalDataloader,
 )
 from nerfactory.datamanagers.datasets import InputDataset
 from nerfactory.datamanagers.pixel_sampler import PixelSampler
@@ -163,7 +164,7 @@ class DataManager(nn.Module):
         raise NotImplementedError
 
     @abstractmethod
-    def next_train(self) -> Tuple:
+    def next_train(self, step: int) -> Tuple:
         """Returns the next batch of data from the train data manager.
 
         This will be a tuple of all the information that this data manager outputs.
@@ -171,13 +172,18 @@ class DataManager(nn.Module):
         raise NotImplementedError
 
     @abstractmethod
-    def next_eval(self) -> Tuple:
+    def next_eval(self, step: int) -> Tuple:
         """Returns the next batch of data from the eval data manager.
 
         This will be a tuple of all the information that this data manager outputs.
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def next_eval_image(self, step: int) -> Tuple:
+        """Returns the next eval image."""
+        raise NotImplementedError
+
     def get_training_callbacks(  # pylint:disable=no-self-use
         self, training_callback_attributes: TrainingCallbackAttributes  # pylint: disable=unused-argument
     ) -> List[TrainingCallback]:
@@ -285,15 +291,22 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
         self.eval_ray_generator = RayGenerator(self.eval_input_dataset.dataset_inputs.cameras.to(self.device))
         # for loading full images
-        self.eval_dataloader = FixedIndicesEvalDataloader(
+        self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(
+            input_dataset=self.eval_input_dataset,
+            device=self.device,
+            num_workers=0 if self.world_size == 1 else self.world_size * 4,
+        )
+        self.eval_dataloader = RandIndicesEvalDataloader(
             input_dataset=self.eval_input_dataset,
-            num_rays_per_chunk=self.config.eval_num_rays_per_chunk,
             image_indices=self.config.eval_image_indices,
             device=self.device,
             num_workers=0 if self.world_size == 1 else self.world_size * 4,
         )
 
-    def next_train(self) -> Tuple[RayBundle, Dict]:
+        # TODO: eval dataloader should be separate from train
+        self.iter_eval_dataloader = iter(self.eval_image_dataloader)
+
+    def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the train dataloader."""
         self.train_count += 1
         image_batch = next(self.iter_train_image_dataloader)
@@ -302,7 +315,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         ray_bundle = self.train_ray_generator(ray_indices)
         return ray_bundle, batch
 
-    def next_eval(self) -> Tuple[RayBundle, Dict]:
+    def next_eval(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the eval dataloader."""
         self.eval_count += 1
         image_batch = next(self.iter_eval_image_dataloader)
@@ -310,3 +323,8 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         ray_indices = batch["indices"]
         ray_bundle = self.eval_ray_generator(ray_indices)
         return ray_bundle, batch
+
+    def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]:
+        for camera_ray_bundle, batch in self.eval_dataloader:
+            image_idx = int(camera_ray_bundle.camera_indices[0, 0])
+            return image_idx, camera_ray_bundle, batch
diff --git a/nerfactory/datamanagers/dataloaders.py b/nerfactory/datamanagers/dataloaders.py
index fd9af9ba5..495d8e186 100644
--- a/nerfactory/datamanagers/dataloaders.py
+++ b/nerfactory/datamanagers/dataloaders.py
@@ -104,20 +104,17 @@ class EvalDataloader(DataLoader):
 
     Args:
         input_dataset: InputDataset to load data from
-        num_rays_per_chunk: Number of camera rays to generate per chunk
         device: Device to load data to
     """
 
     def __init__(
         self,
         input_dataset: InputDataset,
-        num_rays_per_chunk: int,
         device: Union[torch.device, str] = "cpu",
         **kwargs,
     ):
         self.input_dataset = input_dataset
         self.cameras = input_dataset.dataset_inputs.cameras.to(device)
-        self.num_rays_per_chunk = num_rays_per_chunk
         self.device = device
         self.kwargs = kwargs
         super().__init__(dataset=input_dataset)
@@ -159,7 +156,6 @@ class EvalDataloader(DataLoader):
             image_idx: Camera image index
         """
         ray_bundle = self.cameras.generate_rays(camera_indices=image_idx)
-        ray_bundle.num_rays_per_chunk = self.num_rays_per_chunk
         ray_bundle.camera_indices = torch.Tensor([image_idx])[..., None].int()
         batch = self.input_dataset[image_idx]
         batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
@@ -171,7 +167,6 @@ class FixedIndicesEvalDataloader(EvalDataloader):
 
     Args:
         input_dataset: InputDataset to load data from
-        num_rays_per_chunk: Number of camera rays to generate per chunk
         image_indices: List of image indices to load data from. If None, then use all images.
         device: Device to load data to
     """
@@ -179,17 +174,15 @@ class FixedIndicesEvalDataloader(EvalDataloader):
     def __init__(
         self,
         input_dataset: InputDataset,
-        num_rays_per_chunk: int,
         image_indices: Optional[Tuple[int]] = None,
         device: Union[torch.device, str] = "cpu",
         **kwargs,
     ):
-        super().__init__(input_dataset, num_rays_per_chunk, device, **kwargs)
+        super().__init__(input_dataset, device, **kwargs)
         if image_indices is None:
             self.image_indices = list(range(len(input_dataset)))
         else:
             self.image_indices = image_indices
-        self.num_rays_per_chunk = num_rays_per_chunk
         self.count = 0
 
     def __iter__(self):
@@ -200,7 +193,6 @@ class FixedIndicesEvalDataloader(EvalDataloader):
         if self.count < len(self.image_indices):
             image_idx = self.image_indices[self.count]
             ray_bundle, batch = self.get_data_from_image_idx(image_idx)
-            ray_bundle.num_rays_per_chunk = self.num_rays_per_chunk
             self.count += 1
             return ray_bundle, batch
         raise StopIteration
@@ -211,18 +203,16 @@ class RandIndicesEvalDataloader(EvalDataloader):
 
     Args:
         input_dataset: InputDataset to load data from
-        num_rays_per_chunk: Number of camera rays to generate per chunk
         device: Device to load data to
     """
 
     def __init__(
         self,
         input_dataset: InputDataset,
-        num_rays_per_chunk: int,
         device: Union[torch.device, str] = "cpu",
         **kwargs,
     ):
-        super().__init__(input_dataset, num_rays_per_chunk, device, **kwargs)
+        super().__init__(input_dataset, device, **kwargs)
         self.count = 0
 
     def __iter__(self):
diff --git a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py
index 31520d6d3..98659d2ae 100644
--- a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py
+++ b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import logging
+import math
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -62,6 +63,27 @@ class Nerfactory(DataParser):
         No image files found. 
         You should check the file_paths in the transforms.json file to make sure they are correct.
         """
+
+        # filter image_filenames and poses based on train/eval split percentage
+        num_images = len(image_filenames)
+        num_train_images = math.ceil(num_images * self.config.train_split_percentage)
+        num_eval_images = num_images - num_train_images
+        i_all = np.arange(num_images)
+        i_train = np.linspace(
+            0, num_images - 1, num_train_images, dtype=int
+        )  # equally spaced training images starting and ending at 0 and num_images-1
+        i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
+        assert len(i_eval) == num_eval_images
+        if split == "train":
+            indices = i_train
+        elif split == "val" or split == "test":
+            indices = i_eval
+        else:
+            raise ValueError(f"Unknown dataparser split {split}")
+
+        image_filenames = [image_filenames[i] for i in indices]
+        poses = [poses[i] for i in indices]
+
         poses = torch.from_numpy(np.array(poses).astype(np.float32))
         poses = camera_utils.auto_orient_poses(poses, method=self.config.orientation_method)
 
diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py
index ef6bbf1e7..e33fc36e1 100644
--- a/nerfactory/engine/trainer.py
+++ b/nerfactory/engine/trainer.py
@@ -38,6 +38,7 @@ from nerfactory.utils.callbacks import (
     TrainingCallbackLocation,
 )
 from nerfactory.utils.decorators import check_main_thread
+from nerfactory.utils.misc import step_check
 from nerfactory.utils.writer import EventName, TimeWriter
 from nerfactory.viewer.server import viewer_utils
 
@@ -133,12 +134,14 @@ class Trainer:
                             step, location=TrainingCallbackLocation.BEFORE_TRAIN_ITERATION
                         )
 
-                    loss_metric_dict = self.train_iteration(step)
+                    # time the forward pass
+                    loss, loss_dict, metrics_dict = self.train_iteration(step)
 
                     # training callbacks after the training iteration
                     for callback in self.callbacks:
                         callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION)
 
+                    # NOTE(ethan): we should remove this when the viewer is disabled
                     with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t:
                         num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
                         try:
@@ -150,14 +153,36 @@ class Trainer:
                                 "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
                             )
 
-                if step % self.config.logging.steps_per_log == 0:
-                    writer.put_dict(name="Train Metrics and Loss", scalar_dict=loss_metric_dict, step=step)
+                # a batch of train rays
+                if step_check(step, self.config.logging.steps_per_log, run_at_zero=True):
+                    writer.put_scalar(name="Train Loss", scalar=loss, step=step)
+                    writer.put_dict(name="Train Loss Dict", scalar_dict=loss_dict, step=step)
+                    writer.put_dict(name="Train Metrics Dict", scalar_dict=metrics_dict, step=step)
+
+                # a batch of eval rays
+                if step_check(step, self.config.trainer.steps_per_eval_batch, run_at_zero=True):
+                    _, eval_loss_dict, eval_metrics_dict = self.pipeline.get_eval_loss_dict(step=step)
+                    eval_loss = functools.reduce(torch.add, eval_loss_dict.values())
+                    writer.put_scalar(name="Eval Loss", scalar=eval_loss, step=step)
+                    writer.put_dict(name="Eval Loss Dict", scalar_dict=eval_loss_dict, step=step)
+                    writer.put_dict(name="Eval Metrics Dict", scalar_dict=eval_metrics_dict, step=step)
+
+                # one eval image
+                if step_check(step, self.config.trainer.steps_per_eval_image):
+                    metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step)
+                    writer.put_dict(name="Eval Images Metrics", scalar_dict=metrics_dict, step=step)
+                    group = "Eval Images"
+                    for image_name, image in images_dict.items():
+                        writer.put_image(name=group + "/" + image_name, image=image, step=step)
+
+                # all eval images
+                if step_check(step, self.config.trainer.steps_per_eval_all_images):
+                    metrics_dict = self.pipeline.get_average_eval_image_metrics(step=step)
+                    writer.put_dict(name="Eval Images Metrics Dict (all images)", scalar_dict=metrics_dict, step=step)
+
                 if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0:
                     self._save_checkpoint(self.config.trainer.model_dir, step)
-                if step % self.config.trainer.steps_per_test == 0:
-                    metrics_dict = self.pipeline.get_eval_loss_dict(step=step)
-                    writer.put_dict(name="Eval Metrics", scalar_dict=metrics_dict, step=step)
-                    # write out the image dictionary returned too
+
                 self._update_rays_per_sec(train_t, vis_t, step)
                 self._write_out_storage(step)
 
@@ -187,9 +212,9 @@ class Trainer:
         """
         if (
             step % self.config.logging.steps_per_log == 0
-            or (self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0)
-            or step % self.config.trainer.steps_per_test == 0
-            or step == self.config.trainer.max_num_iterations
+            or step % self.config.trainer.steps_per_eval_batch == 0
+            or step % self.config.trainer.steps_per_eval_image == 0
+            or step % self.config.trainer.steps_per_eval_all_images == 0
         ):
             writer.write_out_storage()
 
@@ -259,6 +284,4 @@ class Trainer:
         self.optimizers.scheduler_step_all(step)
 
         # Merging loss and metrics dict into a single output.
-        loss_dict["loss"] = loss
-        loss_dict.update(metrics_dict)
-        return loss_dict
+        return loss, loss_dict, metrics_dict
diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py
index 4f1a931f2..5067c840c 100644
--- a/nerfactory/models/base.py
+++ b/nerfactory/models/base.py
@@ -32,7 +32,7 @@ from nerfactory.datamanagers.structs import SceneBounds
 from nerfactory.fields.density_fields.density_grid import DensityGrid
 from nerfactory.models.modules.scene_colliders import NearFarCollider
 from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes
-from nerfactory.utils.misc import get_masked_dict, is_not_none
+from nerfactory.utils.misc import get_masked_dict
 
 
 class Model(nn.Module):
@@ -192,14 +192,14 @@ class Model(nn.Module):
         Args:
             camera_ray_bundle: ray bundle to calculate outputs over
         """
-        assert is_not_none(camera_ray_bundle.num_rays_per_chunk)
+        num_rays_per_chunk = self.config.eval_num_rays_per_chunk
         image_height, image_width = camera_ray_bundle.origins.shape[:2]
         num_rays = len(camera_ray_bundle)
         outputs = {}
         outputs_lists = defaultdict(list)
-        for i in range(0, num_rays, camera_ray_bundle.num_rays_per_chunk):
+        for i in range(0, num_rays, num_rays_per_chunk):
             start_idx = i
-            end_idx = i + camera_ray_bundle.num_rays_per_chunk
+            end_idx = i + num_rays_per_chunk
             ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
             outputs = self.forward(ray_bundle=ray_bundle)
             for output_name, output in outputs.items():  # type: ignore
@@ -209,7 +209,9 @@ class Model(nn.Module):
         return outputs
 
     @abstractmethod
-    def log_test_image_outputs(self, image_idx, step, batch, outputs) -> Dict[str, float]:
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         """Writes the test image outputs.
         TODO: This shouldn't return a loss
 
diff --git a/nerfactory/models/compound.py b/nerfactory/models/compound.py
index 087171957..8740528f6 100644
--- a/nerfactory/models/compound.py
+++ b/nerfactory/models/compound.py
@@ -28,7 +28,7 @@ TODO:
 
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import torch
 from torch.nn import Parameter
@@ -44,7 +44,7 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames
 from nerfactory.models.base import Model
 from nerfactory.models.modules.ray_sampler import NGPSpacedSampler
 from nerfactory.optimizers.loss import MSELoss
-from nerfactory.utils import colors, misc, visualization, writer
+from nerfactory.utils import colors, misc, visualization
 from nerfactory.utils.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
@@ -164,7 +164,9 @@ class CompoundModel(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         device = self.device
         image = batch["image"].to(device)
         rgb = outputs["rgb"]
@@ -177,10 +179,6 @@ class CompoundModel(Model):
         combined_acc = torch.cat([acc], dim=1)
         combined_depth = torch.cat([depth], dim=1)
 
-        writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
-        writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
-        writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb = torch.moveaxis(rgb, -1, 0)[None, ...]
@@ -189,10 +187,10 @@ class CompoundModel(Model):
         ssim = self.ssim(image, rgb)
         lpips = self.lpips(image, rgb)
 
-        writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(psnr), step=step)
-        writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step)  # type: ignore
-        writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step)
-
-        writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step)
-
-        return {"psnr": float(psnr.item())}
+        metrics_dict = {
+            "psnr": float(psnr.item()),
+            "ssim": float(ssim.item()),
+            "lpips": float(lpips.item()),
+        }
+        images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py
index 3f31b2af5..2a2d30388 100644
--- a/nerfactory/models/instant_ngp.py
+++ b/nerfactory/models/instant_ngp.py
@@ -20,7 +20,7 @@ from __future__ import annotations
 
 import math
 from dataclasses import dataclass, field
-from typing import Dict, List, Literal, Optional, Type
+from typing import Dict, List, Literal, Optional, Tuple, Type
 
 import nerfacc  # pylint: disable=import-error
 import torch
@@ -36,7 +36,7 @@ from nerfactory.fields.instant_ngp_field import field_implementation_to_class
 from nerfactory.fields.modules.field_heads import FieldHeadNames
 from nerfactory.models.base import Model
 from nerfactory.optimizers.loss import MSELoss
-from nerfactory.utils import visualization, writer
+from nerfactory.utils import visualization
 from nerfactory.utils.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
@@ -192,7 +192,10 @@ class NGPModel(Model):
         loss_dict = {"rgb_loss": rgb_loss}
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
+
         image = batch["image"].to(self.device)
         rgb = outputs["rgb"]
         acc = visualization.apply_colormap(outputs["accumulation"])
@@ -207,12 +210,6 @@ class NGPModel(Model):
         combined_depth = torch.cat([depth], dim=1)
         combined_alive_ray_mask = torch.cat([alive_ray_mask], dim=1)
 
-        # TODO: return an image dictionary instead of logging here
-        writer.put_image(name="Eval Images/img", image=combined_rgb, step=step)
-        writer.put_image(name="Eval Images/accumulation", image=combined_acc, step=step)
-        writer.put_image(name="Eval Images/depth", image=combined_depth, step=step)
-        writer.put_image(name="Eval Images/alive_ray_mask", image=combined_alive_ray_mask, step=step)
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb = torch.moveaxis(rgb, -1, 0)[None, ...]
@@ -222,13 +219,14 @@ class NGPModel(Model):
         lpips = self.lpips(image, rgb)
 
         # all of these metrics will be logged as scalars
-        metrics_dict = {
-            "psnr": float(psnr.item()),
-            "image_idx": image_idx,
-            "ssim": float(ssim),  # type: ignore
-            "lpips": float(lpips),
-            "acc_min": float(acc.min().item()),
-            "acc_max": float(acc.max().item()),
-        }
+        metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim), "lpips": float(lpips)}  # type: ignore
         # TODO(ethan): return an image dictionary
-        return metrics_dict
+
+        images_dict = {
+            "img": combined_rgb,
+            "accumulation": combined_acc,
+            "depth": combined_depth,
+            "alive_ray_mask": combined_alive_ray_mask,
+        }
+
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py
index 125bf8a85..51dc6c8f2 100644
--- a/nerfactory/models/mipnerf.py
+++ b/nerfactory/models/mipnerf.py
@@ -17,7 +17,7 @@ Implementation of mip-NeRF.
 """
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import torch
 from torch.nn import Parameter
@@ -38,7 +38,7 @@ from nerfactory.renderers.renderers import (
     DepthRenderer,
     RGBRenderer,
 )
-from nerfactory.utils import colors, misc, visualization, writer
+from nerfactory.utils import colors, misc, visualization
 
 
 class MipNerfModel(Model):
@@ -145,7 +145,9 @@ class MipNerfModel(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
@@ -168,10 +170,6 @@ class MipNerfModel(Model):
         combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)
         combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)
 
-        writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
-        writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
-        writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]
@@ -184,11 +182,12 @@ class MipNerfModel(Model):
         fine_ssim = self.ssim(image, rgb_fine)
         fine_lpips = self.lpips(image, rgb_fine)
 
-        writer.put_scalar(name=f"psnr/val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step)
-        writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step)
-        writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step)  # type: ignore
-        writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step)
-
-        writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step)
-
-        return {"psnr": float(fine_psnr.item())}
+        metrics_dict = {
+            "psnr": float(fine_psnr.item()),
+            "coarse_psnr": float(coarse_psnr.item()),
+            "fine_psnr": float(fine_psnr.item()),
+            "fine_ssim": float(fine_ssim.item()),
+            "fine_lpips": float(fine_lpips.item()),
+        }
+        images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py
index 4ad3c0b2a..89812167d 100644
--- a/nerfactory/models/mipnerf_360.py
+++ b/nerfactory/models/mipnerf_360.py
@@ -18,7 +18,7 @@ Implementation of mip-NeRF.
 
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import torch
 from torch.nn import Parameter
@@ -41,7 +41,7 @@ from nerfactory.renderers.renderers import (
     DepthRenderer,
     RGBRenderer,
 )
-from nerfactory.utils import colors, misc, visualization, writer
+from nerfactory.utils import colors, misc, visualization
 
 
 class MipNerf360Model(Model):
@@ -160,7 +160,9 @@ class MipNerf360Model(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
@@ -183,10 +185,6 @@ class MipNerf360Model(Model):
         combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)
         combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)
 
-        writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
-        writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
-        writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]
@@ -197,17 +195,14 @@ class MipNerf360Model(Model):
         fine_ssim = self.ssim(image, rgb_fine)
         fine_lpips = self.lpips(image, rgb_fine)
 
-        writer.put_scalar(name=f"psnr/val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step)
-        writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step)
-        writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step)  # type: ignore
-        writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step)
-        writer.put_scalar(
-            name=f"ray_loss_coarse/val_{image_idx}", scalar=float(torch.mean(outputs["ray_loss_coarse"])), step=step
-        )
-        writer.put_scalar(
-            name=f"ray_loss_fine/val_{image_idx}", scalar=float(torch.mean(outputs["ray_loss_fine"])), step=step
-        )
-
-        writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step)
-
-        return {"psnr": float(fine_psnr.item())}
+        metrics_dict = {
+            "psnr": float(fine_psnr.item()),
+            "coarse_psnr": float(coarse_psnr),
+            "fine_psnr": float(fine_psnr),
+            "fine_ssim": float(fine_ssim),
+            "fine_lpips": float(fine_lpips),
+            "ray_loss_coarse": float(torch.mean(outputs["ray_loss_coarse"])),
+            "ray_loss_fine": float(torch.mean(outputs["ray_loss_fine"])),
+        }
+        images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py
index 3255c6c3a..1105cf57f 100644
--- a/nerfactory/models/modules/ray_sampler.py
+++ b/nerfactory/models/modules/ray_sampler.py
@@ -19,7 +19,7 @@ Collection of sampling strategies
 from abc import abstractmethod
 from typing import Callable, Optional, Tuple
 
-import nerfacc
+import nerfacc  # pylint: disable=import-error
 import torch
 from torch import nn
 from torchtyping import TensorType
diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py
index 75cfd8f65..b6d49e19b 100644
--- a/nerfactory/models/nerfw.py
+++ b/nerfactory/models/nerfw.py
@@ -18,6 +18,8 @@ NeRF-W (NeRF in the wild) implementation.
 
 from __future__ import annotations
 
+from typing import Dict, Tuple
+
 import torch
 from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.functional import structural_similarity_index_measure
@@ -39,7 +41,7 @@ from nerfactory.renderers.renderers import (
     RGBRenderer,
     UncertaintyRenderer,
 )
-from nerfactory.utils import colors, misc, visualization, writer
+from nerfactory.utils import colors, misc, visualization
 
 
 class NerfWModel(Model):
@@ -214,7 +216,9 @@ class NerfWModel(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
@@ -234,14 +238,13 @@ class NerfWModel(Model):
         row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2)
         combined_image = torch.cat([row0, row1, row2], dim=-3)
 
-        writer.put_image(name=f"img/image_idx_{image_idx}-nerfw", image=combined_image, step=step)
-
-        if "mask" in batch:
-            mask = batch["mask"].repeat(1, 1, 3)
-            writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step)
-
         # this doesn't really make sense but do it anyway
         fine_psnr = self.psnr(image, rgb_fine)
-        return {
+        metrics_dict = {
             "psnr": float(fine_psnr.item()),
         }
+        images_dict = {"img": combined_image}
+        if "mask" in batch:
+            mask = batch["mask"].repeat(1, 1, 3)
+            images_dict["mask"] = mask
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py
index 6909d9c7a..b90a12cdd 100644
--- a/nerfactory/models/semantic_nerf.py
+++ b/nerfactory/models/semantic_nerf.py
@@ -17,6 +17,8 @@ Semantic NeRF implementation.
 """
 from __future__ import annotations
 
+from typing import Dict, Tuple
+
 import torch
 from torch import nn
 
@@ -30,7 +32,7 @@ from nerfactory.fields.semantic_nerf_field import SemanticNerfField
 from nerfactory.models.modules.scene_colliders import AABBBoxCollider
 from nerfactory.models.vanilla_nerf import NeRFModel
 from nerfactory.renderers.renderers import SemanticRenderer
-from nerfactory.utils import misc, writer
+from nerfactory.utils import misc
 
 
 class SemanticNerfModel(NeRFModel):
@@ -128,12 +130,16 @@ class SemanticNerfModel(NeRFModel):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        super().log_test_image_outputs(image_idx, step, batch, outputs)
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
+        metrics_dict, images_dict = super().get_image_metrics_and_images(outputs, batch)
         semantic_logits = outputs["semantic_fine"]
         semantic_labels = torch.argmax(torch.nn.functional.softmax(semantic_logits, dim=-1), dim=-1)  # type: ignore
         semantic_labels_image = self.stuff_colors[semantic_labels]
-        writer.put_image(name=f"semantics/image_idx_{image_idx}", image=semantic_labels_image, step=step)
 
         mask = batch["mask"].repeat(1, 1, 3)
-        writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step)
+
+        images_dict["semantics"] = semantic_labels_image
+        images_dict["mask"] = mask
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py
index 1d4fcb9d2..c739e2d59 100644
--- a/nerfactory/models/tensorf.py
+++ b/nerfactory/models/tensorf.py
@@ -18,7 +18,7 @@ TensorRF implementation.
 
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import numpy as np
 import torch
@@ -41,7 +41,7 @@ from nerfactory.renderers.renderers import (
     DepthRenderer,
     RGBRenderer,
 )
-from nerfactory.utils import colors, misc, visualization, writer
+from nerfactory.utils import colors, misc, visualization
 from nerfactory.utils.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
@@ -195,7 +195,9 @@ class TensoRFModel(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb"].device)
         rgb = outputs["rgb"]
         acc = visualization.apply_colormap(outputs["accumulation"])
@@ -208,10 +210,6 @@ class TensoRFModel(Model):
 
         combined_rgb = torch.cat([image, rgb], dim=1)
 
-        writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
-        writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=acc, step=step)
-        writer.put_image(name=f"depth/image_idx_{image_idx}", image=depth, step=step)
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb = torch.moveaxis(rgb, -1, 0)[None, ...]
@@ -220,12 +218,10 @@ class TensoRFModel(Model):
         ssim = self.ssim(image, rgb)
         lpips = self.lpips(image, rgb)
 
-        writer.put_scalar(name=f"psnr/val_{image_idx}", scalar=float(psnr), step=step)
-        writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step)  # type: ignore
-        writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step)
-
-        writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step)
-
-        return {
+        metrics_dict = {
             "psnr": float(psnr.item()),
+            "ssim": float(ssim.item()),
+            "lpips": float(lpips.item()),
         }
+        images_dict = {"img": combined_rgb, "accumulation": acc, "depth": depth}
+        return metrics_dict, images_dict
diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py
index b62f1ab7c..4693c797e 100644
--- a/nerfactory/models/vanilla_nerf.py
+++ b/nerfactory/models/vanilla_nerf.py
@@ -18,7 +18,7 @@ Implementation of vanilla nerf.
 
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import torch
 from torch.nn import Parameter
@@ -39,7 +39,7 @@ from nerfactory.renderers.renderers import (
     DepthRenderer,
     RGBRenderer,
 )
-from nerfactory.utils import colors, misc, visualization, writer
+from nerfactory.utils import colors, misc, visualization
 from nerfactory.utils.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
@@ -178,7 +178,9 @@ class NeRFModel(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
@@ -201,10 +203,6 @@ class NeRFModel(Model):
         combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)
         combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)
 
-        writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
-        writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
-        writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
-
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]
         rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...]
@@ -215,11 +213,14 @@ class NeRFModel(Model):
         fine_ssim = self.ssim(image, rgb_fine)
         fine_lpips = self.lpips(image, rgb_fine)
 
-        writer.put_scalar(name=f"psnr/val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step)
-        writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step)
-        writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step)  # type: ignore
-        writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step)
-
-        writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step)
-
-        return {"psnr": float(fine_psnr.item())}
+        metrics_dict = {
+            "psnr": float(fine_psnr.item()),
+            "coarse_psnr": float(coarse_psnr),
+            "fine_psnr": float(fine_psnr),
+            "fine_ssim": float(fine_ssim),
+            "fine_lpips": float(fine_lpips),
+            "ray_loss_coarse": float(torch.mean(outputs["ray_loss_coarse"])),
+            "ray_loss_fine": float(torch.mean(outputs["ray_loss_fine"])),
+        }
+        images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
+        return metrics_dict, images_dict
diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py
index 19ca0091d..66addd11b 100644
--- a/nerfactory/pipelines/base.py
+++ b/nerfactory/pipelines/base.py
@@ -21,6 +21,7 @@ import typing
 from abc import abstractmethod
 from typing import Any, Dict, List, Optional
 
+import torch
 import torch.distributed as dist
 from torch import nn
 from torch.nn import Parameter
@@ -34,6 +35,15 @@ from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttribu
 from nerfactory.utils.writer import EventName, TimeWriter
 
 
+def module_wrapper(module: nn.Module) -> nn.Module:
+    """
+    If DDP, then return the .module. Otherwise, return the model.
+    """
+    if isinstance(module, DDP):
+        return module.module
+    return module
+
+
 class Pipeline(nn.Module):
     """The intent of this class is to provide a higher level interface for the Model
     that will be easy to use for our Trainer class.
@@ -103,7 +113,7 @@ class Pipeline(nn.Module):
     @property
     def device(self):
         """Returns the device that the model is on."""
-        return self.model.module.device if self.world_size > 1 else self.model.device
+        return module_wrapper(self.model).device
 
     @abstractmethod
     @profiler.time_function
@@ -116,8 +126,8 @@ class Pipeline(nn.Module):
             step: current iteration step to update sampler if using DDP (distributed)
         """
         if self.world_size > 1:
-            self.datamanager.sampler.set_epoch(step)
-        ray_bundle, batch = self.datamanager.next_train()
+            self.datamanager.train_image_dataloader.sampler.set_epoch(step)
+        ray_bundle, batch = self.datamanager.next_train(step)
         model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch)
         return model_outputs, loss_dict, metrics_dict
 
@@ -131,25 +141,54 @@ class Pipeline(nn.Module):
             step: current iteration step
         """
         self.eval()
-        # NOTE(ethan): next_eval() is not being used right now
-        assert self.datamanager.eval_dataloader is not None
-        for camera_ray_bundle, batch in self.datamanager.eval_dataloader:
-            with TimeWriter(None, None, write=False) as test_t:
-                assert camera_ray_bundle.camera_indices is not None
-                image_idx = int(camera_ray_bundle.camera_indices[0, 0])
-                if self.world_size > 1:
-                    outputs = self.model.module.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
-                    metrics_dict = self.model.module.log_test_image_outputs(image_idx, step, batch, outputs)
-                else:
-                    outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
-                    metrics_dict = self.model.log_test_image_outputs(image_idx, step, batch, outputs)
-            if step is not None:
-                writer.put_time(
-                    name=EventName.TEST_RAYS_PER_SEC,
-                    duration=len(camera_ray_bundle) / test_t.duration,
-                    step=step,
-                    avg_over_steps=True,
-                )
+        if self.world_size > 1:
+            self.datamanager.eval_image_dataloader.sampler.set_epoch(step)
+        ray_bundle, batch = self.datamanager.next_eval(step)
+        model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch)
+        self.train()
+        return model_outputs, loss_dict, metrics_dict
+
+    @abstractmethod
+    @profiler.time_function
+    def get_eval_image_metrics_and_images(self, step: Optional[int] = None):
+        """This function gets your evaluation loss dict. It needs to get the data
+        from the DataManager and feed it to the model's forward function
+
+        Args:
+            step: current iteration step
+        """
+        self.eval()
+        image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step)
+        with TimeWriter(None, None, write=False) as test_t:
+            outputs = module_wrapper(self.model).get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+            metrics_dict, images_dict = module_wrapper(self.model).get_image_metrics_and_images(outputs, batch)
+        writer.put_time(
+            name=EventName.TEST_RAYS_PER_SEC,
+            duration=len(camera_ray_bundle) / test_t.duration,
+            step=step,
+            avg_over_steps=True,
+        )
+        metrics_dict["image_idx"] = image_idx
+        self.train()
+        return metrics_dict, images_dict
+
+    @abstractmethod
+    @profiler.time_function
+    def get_average_eval_image_metrics(self, step: Optional[int] = None):
+        """Iterate over all the images in the eval dataset and get the average."""
+        self.eval()
+        metrics_dict_list = []
+        # TODO: add something like tqdm but so that it doesn't interfere with logging
+        for camera_ray_bundle, batch in self.datamanager.fixed_indices_eval_dataloader:
+            outputs = module_wrapper(self.model).get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+            metrics_dict, _ = module_wrapper(self.model).get_image_metrics_and_images(outputs, batch)
+            metrics_dict_list.append(metrics_dict)
+        # average the metrics list
+        metrics_dict = {}
+        for key in metrics_dict_list[0].keys():
+            metrics_dict[key] = float(
+                torch.mean(torch.tensor([metrics_dict[key] for metrics_dict in metrics_dict_list]))
+            )
         self.train()
         return metrics_dict
 
@@ -171,10 +210,7 @@ class Pipeline(nn.Module):
     ) -> List[TrainingCallback]:
         """Returns the training callbacks from both the Dataloader and the Model."""
         datamanager_callbacks = self.datamanager.get_training_callbacks(training_callback_attributes)
-        if self.world_size > 1:
-            model_callbacks = self.model.module.get_training_callbacks(training_callback_attributes)
-        else:
-            model_callbacks = self.model.get_training_callbacks(training_callback_attributes)
+        model_callbacks = module_wrapper(self.model).get_training_callbacks(training_callback_attributes)
         callbacks = datamanager_callbacks + model_callbacks
         return callbacks
 
@@ -185,9 +221,6 @@ class Pipeline(nn.Module):
             A list of dictionaries containing the pipeline's param groups.
         """
         datamanager_params = self.datamanager.get_param_groups()
-        if self.world_size > 1:
-            model_params = self.model.module.get_param_groups()
-        else:
-            model_params = self.model.get_param_groups()
+        model_params = module_wrapper(self.model).get_param_groups()
         # TODO(ethan): assert that key names don't overlap
         return {**datamanager_params, **model_params}
diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py
index 297c5d536..f18df1706 100644
--- a/nerfactory/utils/misc.py
+++ b/nerfactory/utils/misc.py
@@ -159,3 +159,8 @@ def scale_dict(dictionary: Dict[Any, Any], coefficients: Dict[str, float]) -> Di
         if key in coefficients:
             dictionary[key] *= coefficients[key]
     return dictionary
+
+
+def step_check(step, step_size, run_at_zero=False) -> bool:
+    """Returns true based on current step and step interval."""
+    return (run_at_zero or step != 0) and step % step_size == 0
diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py
index 86d445564..488dac42c 100644
--- a/nerfactory/utils/writer.py
+++ b/nerfactory/utils/writer.py
@@ -22,7 +22,7 @@ import logging
 from abc import abstractmethod
 from pathlib import Path
 from time import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import imageio
 import numpy as np
@@ -79,7 +79,7 @@ def put_image(name, image: TensorType["H", "W", "C"], step: int):
 
 
 @check_main_thread
-def put_scalar(name: str, scalar: float, step: int):
+def put_scalar(name: str, scalar: Any, step: int):
     """Setter function to place scalars into the queue to be written out
 
     Args:
@@ -209,7 +209,7 @@ class Writer:
         raise NotImplementedError
 
     @abstractmethod
-    def write_scalar(self, name: str, scalar: float, step: int) -> None:
+    def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None:
         """Required method to write a single scalar value to the logger
 
         Args:
@@ -272,7 +272,7 @@ class WandbWriter(Writer):
         image = torch.permute(image, (2, 0, 1))
         wandb.log({name: wandb.Image(image)}, step=step)
 
-    def write_scalar(self, name: str, scalar: float, step: int) -> None:
+    def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None:
         wandb.log({name: scalar}, step=step)
 
     def write_config(self, name: str, config_dict: Dict[str, Any], step: int):
@@ -298,7 +298,7 @@ class TensorboardWriter(Writer):
         image = to8b(image)
         self.tb_writer.add_image(name, image, step, dataformats="HWC")
 
-    def write_scalar(self, name: str, scalar: float, step: int) -> None:
+    def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None:
         self.tb_writer.add_scalar(name, scalar, step)
 
     def write_config(self, name: str, config_dict: Dict[str, Any], step: int):  # pylint: disable=unused-argument
diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py
index 130772430..2b3d353f8 100644
--- a/nerfactory/viewer/server/viewer_utils.py
+++ b/nerfactory/viewer/server/viewer_utils.py
@@ -616,7 +616,6 @@ class VisualizerState:
         camera = camera.to(graph.device)
 
         camera_ray_bundle = camera.generate_rays(camera_indices=0)
-        camera_ray_bundle.num_rays_per_chunk = self.config.num_rays_per_chunk
 
         graph.eval()
 
diff --git a/scripts/run_eval.py b/scripts/run_eval.py
index 84e1fea7e..2cc34fc55 100755
--- a/scripts/run_eval.py
+++ b/scripts/run_eval.py
@@ -100,7 +100,6 @@ def _render_trajectory_video(
     output_filename: Path,
     rendered_output_name: str,
     rendered_resolution_scaling_factor: float = 1.0,
-    num_rays_per_chunk: int = 4096,
     seconds: float = 5.0,
 ) -> None:
     """Helper function to create a video of the spiral trajectory.
@@ -112,7 +111,6 @@ def _render_trajectory_video(
         rendered_output_name: Name of the renderer output to use.
         rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution.
             Defaults to 1.0.
-        num_rays_per_chunk: Number of rays to use per chunk. Defaults to 4096.
         seconds: Number for the output video. Defaults to 5.0.
     """
     console.print("[bold green]Creating trajectory video")
@@ -120,7 +118,6 @@ def _render_trajectory_video(
     cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
     for camera_idx in track(range(cameras.size), description=":movie_camera: Rendering :movie_camera:"):
         camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx).to(pipeline.device)
-        camera_ray_bundle.num_rays_per_chunk = num_rays_per_chunk
         with torch.no_grad():
             outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
         image = outputs[rendered_output_name].cpu().numpy()
@@ -217,7 +214,7 @@ class RenderTrajectory:
 
     def main(self) -> None:
         """Main function."""
-        config, pipeline, _ = _eval_setup(self.load_config)
+        _, pipeline, _ = _eval_setup(self.load_config)
 
         seconds = self.seconds
 
@@ -245,7 +242,6 @@ class RenderTrajectory:
             output_filename=self.output_path,
             rendered_output_name=self.rendered_output_name,
             rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
-            num_rays_per_chunk=config.pipeline.datamanager.eval_num_rays_per_chunk,
             seconds=seconds,
         )