diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py
index 6c1722215..2cf823d46 100644
--- a/nerfactory/datamanagers/base.py
+++ b/nerfactory/datamanagers/base.py
@@ -88,6 +88,8 @@ class DataManager(nn.Module):
 
     train_input_dataset: Optional[InputDataset] = None
     eval_input_dataset: Optional[InputDataset] = None
+    train_sampler: Optional[DistributedSampler] = None
+    eval_sampler: Optional[DistributedSampler] = None
 
     def __init__(self):
         """Constructor for the DataManager class.
@@ -326,5 +328,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
 
     def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]:
         for camera_ray_bundle, batch in self.eval_dataloader:
+            assert camera_ray_bundle.camera_indices is not None
             image_idx = int(camera_ray_bundle.camera_indices[0, 0])
             return image_idx, camera_ray_bundle, batch
+        raise ValueError("No more eval images")
diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py
index 6b7401a5f..10183f82f 100644
--- a/nerfactory/pipelines/base.py
+++ b/nerfactory/pipelines/base.py
@@ -117,7 +117,7 @@ class Pipeline(nn.Module):
 
     @abstractmethod
     @profiler.time_function
-    def get_train_loss_dict(self, step: Optional[int] = None):
+    def get_train_loss_dict(self, step: int):
         """This function gets your training loss dict. This will be responsible for
         getting the next batch of data from the DataManager and interfacing with the
         Model class, feeding the data to the model's forward function.
@@ -125,15 +125,16 @@ class Pipeline(nn.Module):
         Args:
             step: current iteration step to update sampler if using DDP (distributed)
         """
-        if self.world_size > 1:
-            self.datamanager.train_image_dataloader.sampler.set_epoch(step)
+        if self.world_size > 1 and step:
+            assert self.datamanager.train_sampler is not None
+            self.datamanager.train_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
 
     @abstractmethod
     @profiler.time_function
-    def get_eval_loss_dict(self, step: Optional[int] = None):
+    def get_eval_loss_dict(self, step: int):
         """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
 
@@ -142,7 +143,8 @@ class Pipeline(nn.Module):
         """
         self.eval()
         if self.world_size > 1:
-            self.datamanager.eval_image_dataloader.sampler.set_epoch(step)
+            assert self.datamanager.eval_sampler is not None
+            self.datamanager.eval_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()
@@ -150,7 +152,7 @@ class Pipeline(nn.Module):
 
     @abstractmethod
     @profiler.time_function
-    def get_eval_image_metrics_and_images(self, step: Optional[int] = None):
+    def get_eval_image_metrics_and_images(self, step: int):
         """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