throw exception for inappropriate usage

ddp_train_nerf.py

@@ -121,6 +121,11 @@ def sample_pdf(bins, weights, N_samples, det=False):
 def render_single_image(rank, world_size, models, ray_sampler, chunk_size):
     ##### parallel rendering of a single image
     ray_batch = ray_sampler.get_all()
+    
+    if (ray_batch['ray_d'].shape[0] // world_size) * world_size != ray_batch['ray_d'].shape[0]:
+        raise Exception('Number of pixels in the image is not divisible by the number of GPUs!\n\t# pixels: {}\n\t# GPUs: {}'.format(ray_batch['ray_d'].shape[0],
+                                                                                                                                     world_size))
+    
     # split into ranks; make sure different processes don't overlap
     rank_split_sizes = [ray_batch['ray_d'].shape[0] // world_size, ] * world_size
     rank_split_sizes[-1] = ray_batch['ray_d'].shape[0] - sum(rank_split_sizes[:-1])