move images to device properly

Author: kerrj

nerfstudio/data/utils/dataloaders.py

@@ -631,6 +631,12 @@ def __iter__(self):
                 camera, data = self.custom_view_processor(camera, data)
 
             i += 1
+            camera = camera.to(self.device)
+            data['image'] = data['image'].to(self.device)
+            if 'mask' in data:
+                data['mask'] = data['mask'].to(self.device)
+            if 'depth_image' in data:
+                data['depth_image'] = data['depth_image'].to(self.device)
             yield camera, data