Move only per-chunk data to the GPU to support rendering large images (#2481)

* Move only per-chunk data to the GPU to support large images

* make type checking happy

* Use isinstance check

torch.is_tensor() doesn't have a TypeGuard annotation

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>

nerfstudio/models/base_model.py

@@ -168,6 +168,7 @@ class Model(nn.Module):
         Args:
             camera_ray_bundle: ray bundle to calculate outputs over
         """
+        input_device = camera_ray_bundle.directions.device
         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)
@@ -176,12 +177,15 @@ class Model(nn.Module):
             start_idx = i
             end_idx = i + num_rays_per_chunk
             ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
+            # move the chunk inputs to the model device
+            ray_bundle = ray_bundle.to(self.device)
             outputs = self.forward(ray_bundle=ray_bundle)
             for output_name, output in outputs.items():  # type: ignore
-                if not torch.is_tensor(output):
+                if not isinstance(output, torch.Tensor):
                     # TODO: handle lists of tensors as well
                     continue
-                outputs_lists[output_name].append(output)
+                # move the chunk outputs from the model device back to the device of the inputs.
+                outputs_lists[output_name].append(output.to(input_device))
         outputs = {}
         for output_name, outputs_list in outputs_lists.items():
             outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1)  # type: ignore