Fix datamanagers' get_rays_per_batch functions in case of using DynamicBatchPipeline (#2522)

nerfstudio/data/datamanagers/base_datamanager.py

@@ -556,9 +556,13 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
         raise ValueError("No more eval images")
 
     def get_train_rays_per_batch(self) -> int:
+        if self.train_pixel_sampler is not None:
+            return self.train_pixel_sampler.num_rays_per_batch
         return self.config.train_num_rays_per_batch
 
     def get_eval_rays_per_batch(self) -> int:
+        if self.eval_pixel_sampler is not None:
+            return self.eval_pixel_sampler.num_rays_per_batch
         return self.config.eval_num_rays_per_batch
 
     def get_datapath(self) -> Path:

nerfstudio/data/datamanagers/parallel_datamanager.py

@@ -303,10 +303,14 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
 
     def get_train_rays_per_batch(self) -> int:
         """Returns the number of rays per batch for training."""
+        if self.train_pixel_sampler is not None:
+            return self.train_pixel_sampler.num_rays_per_batch
         return self.config.train_num_rays_per_batch
 
     def get_eval_rays_per_batch(self) -> int:
         """Returns the number of rays per batch for evaluation."""
+        if self.eval_pixel_sampler is not None:
+            return self.eval_pixel_sampler.num_rays_per_batch
         return self.config.eval_num_rays_per_batch
 
     def get_datapath(self) -> Path: