cpu compatibility

scene/gaussian_predictor.py

@@ -746,7 +746,8 @@ class GaussianSplatPredictor(nn.Module):
         # Pos prediction is in camera space - compute the positions in the world space
         pos = self.flatten_vector(pos)
         pos = torch.cat([pos, 
-                         torch.ones((pos.shape[0], pos.shape[1], 1), device="cuda", dtype=torch.float32)
+                         torch.ones((pos.shape[0], pos.shape[1], 1), 
+                                    device=pos.device, dtype=torch.float32)
                          ], dim=2)
         pos = torch.bmm(pos, source_cameras_view_to_world)
         pos = pos[:, :, :3] / (pos[:, :, 3:] + 1e-10)
@@ -781,7 +782,8 @@ class GaussianSplatPredictor(nn.Module):
             out_dict["features_rest"] = torch.zeros((out_dict["features_dc"].shape[0], 
                                                      out_dict["features_dc"].shape[1], 
                                                      (self.cfg.model.max_sh_degree + 1) ** 2 - 1,
-                                                     3), dtype=out_dict["features_dc"].dtype, device="cuda")
+                                                     3), dtype=out_dict["features_dc"].dtype, 
+                                                     device=out_dict["xyz"].device)
 
         out_dict = self.multi_view_union(out_dict, B, N_views)
         out_dict = self.make_contiguous(out_dict)