check device in run_depth

app.py

@@ -239,7 +239,7 @@ def run_depth_estimation(
     ).to(dtype=torch.float32, device=device)
 
     dpt_pack = {
-        'c2w': c2w.to('cpu'),
+        'c2w': c2w,
         'K': K,
         'depth_mask': depth_mask,
         'depth': depth
@@ -256,14 +256,14 @@ def run_depth_estimation(
         else:
             instance_mask = (seg_image.reshape(-1, 3) == instance_label).all(axis=-1).reshape(H, W)
             mask_pack.append(instance_mask)
-    fg_mask = torch.from_numpy(np.stack(mask_pack).any(axis=0)).to('cpu')
+    fg_mask = torch.from_numpy(np.stack(mask_pack).any(axis=0))
 
     scene_est_depth_pts, scene_est_depth_pts_colors = \
-        project2ply(depth_mask, depth, input_image, K, c2w)
+        project2ply(depth_mask.to(device), depth.to(device), input_image.to(device), K.to(device), c2w.to(device))
     save_ply_path = os.path.join(work_space, "scene_pcd.glb")
 
     fg_depth_pts, _ = \
-        project2ply(fg_mask, depth, input_image, K, c2w)
+        project2ply(fg_mask.to(device), depth.to(device), input_image.to(device), K.to(device), c2w.to(device))
     _, trans, scale = normalize_vertices(fg_depth_pts)
 
     if trans.shape[0] == 1: