Keep depth_image on CPU (#3449)

Co-authored-by: J.Y. <132313008+jb-ye@users.noreply.github.com>

nerfstudio/data/datasets/depth_dataset.py

@@ -41,6 +41,8 @@ class DepthDataset(InputDataset):
         scale_factor: The scaling factor for the dataparser outputs.
     """
 
+    exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth_image"]
+
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
         # if there are no depth images than we want to generate them all with zoe depth