fix CI pyright issues

nerfstudio-project · Dec 7, 2023 · 0ee50a8 · 0ee50a8
1 parent 4719567
commit 0ee50a8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
@@ -284,7 +284,7 @@ def train(self) -> None:
                     # (https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management)
                     # for more details about GPU memory management.
                     writer.put_scalar(
-                        name="GPU Memory (MB)", scalar=torch.cuda.max_memory_allocated() / (1024 ** 2), step=step
+                        name="GPU Memory (MB)", scalar=torch.cuda.max_memory_allocated() / (1024**2), step=step
                     )
 
                 # Do not perform evaluation if there are no validation images

diff --git a/nerfstudio/scripts/train.py b/nerfstudio/scripts/train.py
@@ -109,7 +109,7 @@ def _distributed_worker(
     dist_url: str,
     config: TrainerConfig,
     timeout: timedelta = DEFAULT_TIMEOUT,
-    device_type: Literal["cpu", "cuda", "mps"] = "cuda",
+    device_type: Literal["cpu", "cuda", "mps", "xpu"] = "cuda",
 ) -> Any:
     """Spawned distributed worker that handles the initialization of process group and handles the
        training process on multiple processes.
@@ -165,7 +165,7 @@ def launch(
     dist_url: str = "auto",
     config: Optional[TrainerConfig] = None,
     timeout: timedelta = DEFAULT_TIMEOUT,
-    device_type: Literal["cpu", "cuda", "mps"] = "cuda",
+    device_type: Literal["cpu", "cuda", "mps", "xpu"] = "cuda",
 ) -> None:
     """Function that spawns multiple processes to call on main_func