Literaly my last gueess and then I go do smth else in my life

demo/gs_train.py

@@ -46,7 +46,7 @@ class ModelParams:
     images: str = "images"
     resolution: int = -1
     white_background: bool = True
-    data_device: str = DEVICE
+    data_device: str = "cuda"
     eval: bool = False
 
 @dataclass
@@ -155,7 +155,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=DEVICE)
+    background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
 
     iter_start = torch.cuda.Event(enable_timing = True)
     iter_end = torch.cuda.Event(enable_timing = True)
@@ -251,7 +251,7 @@ def train(
         iteration = scene.loaded_iter
 
         bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
-        background = torch.tensor(bg_color, dtype=torch.float32, device=DEVICE)
+        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
 
         model_path = dataset.model_path
         name = "render"