Check allocation problem

demo/gs_train.py

@@ -7,8 +7,11 @@ from tqdm.auto import tqdm
 import gradio as gr
 import importlib.util
 from dataclasses import dataclass, field
+from demo_globals import DEVICE
 
 import spaces
+from simple_knn._C import distCUDA2
+
 
 
 @dataclass
@@ -45,7 +48,7 @@ class ModelParams:
     images: str = "images"
     resolution: int = -1
     white_background: bool = True
-    data_device: str = "cuda"
+    data_device: str = DEVICE
     eval: bool = False
 
 @dataclass
@@ -60,7 +63,7 @@ class TrainingArgs:
     checkpoint_iterations: list[int] = field(default_factory=lambda: [7_000, 15_000, 30_000])
     start_checkpoint: str = None
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=20)
 def train(
     data_source_path, sh_degree, model_path, images, resolution, white_background, data_device, eval,
     convert_SHs_python, compute_cov3D_python, debug,
@@ -138,6 +141,10 @@ def train(
     checkpoint_iterations = args.checkpoint_iterations 
     debug_from = args.debug_from
 
+    dist2 = torch.clamp_min(distCUDA2(torch.randn((90804, 3)).float().cuda()), 0.0000001)
+    print("dist2.shape: ", dist2.shape)
+
+
     tb_writer = prepare_output_and_logger(dataset)
     
     gaussians = GaussianModel(dataset.sh_degree)
@@ -145,7 +152,7 @@ def train(
     gaussians.training_setup(opt)
 
     bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
-    background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
+    background = torch.tensor(bg_color, dtype=torch.float32, device=DEVICE)
 
     iter_start = torch.cuda.Event(enable_timing = True)
     iter_end = torch.cuda.Event(enable_timing = True)
@@ -174,7 +181,7 @@ def train(
         # Render
         if (iteration - 1) == debug_from:
             pipe.debug = True
-        bg = torch.rand((3), device="cuda") if opt.random_background else background
+        bg = torch.rand((3), device=DEVICE) if opt.random_background else background
         
         render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
         image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]