fix cpu ram following graphdeco-inria#667

Author: nnmhuy

utils/camera_utils.py

@@ -39,7 +39,8 @@ def loadCam(args, id, cam_info, resolution_scale):
         scale = float(global_down) * float(resolution_scale)
         resolution = (int(orig_w / scale), int(orig_h / scale))
 
-    resized_image_rgb = PILtoTorch(cam_info.image, resolution)
+    # resized_image_rgb = PILtoTorch(cam_info.image, resolution)
+    resized_image_rgb = PILtoTorch(cam_info.image, resolution, args.data_device)
 
     gt_image = resized_image_rgb[:3, ...]
     loaded_mask = None

utils/general_utils.py

@@ -18,9 +18,11 @@
 def inverse_sigmoid(x):
     return torch.log(x/(1-x))
 
-def PILtoTorch(pil_image, resolution):
+def PILtoTorch(pil_image, resolution, data_device="cuda"):
     resized_image_PIL = pil_image.resize(resolution)
-    resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
+    # resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
+    resized_image = torch.from_numpy(np.array(resized_image_PIL)).to(data_device) / 255.0
+
     if len(resized_image.shape) == 3:
         return resized_image.permute(2, 0, 1)
     else: