diff --git a/nerfactory/dataloaders/base.py b/nerfactory/dataloaders/base.py
index fa2fb26d7..15c3b8185 100644
--- a/nerfactory/dataloaders/base.py
+++ b/nerfactory/dataloaders/base.py
@@ -262,13 +262,18 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             raise ValueError(f"Unknown image dataset type {self.config.image_dataset_type}")
 
         if self.world_size > 1:
+            self.sampler = DistributedSampler(
+                self.train_image_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
+            )
             self.train_image_dataloader = CacheImageDataloader(
                 self.train_image_dataset,
                 num_images_to_sample_from=self.config.train_num_images_to_sample_from,
                 device=self.device,
                 num_workers=self.world_size * 4,
                 pin_memory=True,
+                sampler=self.sampler,
             )  # TODO(ethan): pass this in
+
         else:
             self.train_image_dataloader = CacheImageDataloader(
                 self.train_image_dataset,
@@ -291,11 +296,6 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         elif self.config.image_dataset_type == "panoptic":
             self.eval_image_dataset = PanopticImageDataset(**self.eval_datasetinputs.as_dict())
 
-        if self.world_size > 1:
-            self.sampler = DistributedSampler(
-                self.train_image_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
-            )
-
         self.eval_dataloader = FixedIndicesEvalDataloader(
             image_dataset=self.eval_image_dataset,
             cameras=self.eval_cameras,
diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py
index 9b1dac099..038866049 100644
--- a/nerfactory/pipelines/base.py
+++ b/nerfactory/pipelines/base.py
@@ -145,7 +145,10 @@ class Pipeline(nn.Module):
     ) -> List[TrainingCallback]:
         """Returns the training callbacks from both the Dataloader and the Model."""
         data_manager_callbacks = self.data_manager.get_training_callbacks(training_callback_attributes)
-        model_callbacks = self.model.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)
         callbacks = data_manager_callbacks + model_callbacks
         return callbacks