Drop set_cuda_device

nerfstudio-project · Aug 14, 2024 · 08e8361 · 08e8361
1 parent 7d99a95
commit 08e8361
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/nerfstudio/scripts/train.py b/nerfstudio/scripts/train.py
@@ -95,8 +95,6 @@ def train_loop(local_rank: int, world_size: int, config: TrainerConfig, global_r
         config: config file specifying training regimen
     """
     _set_random_seed(config.machine.seed + global_rank)
-    if torch.cuda.is_available():
-        torch.cuda.set_device(local_rank)
     trainer = config.setup(local_rank=local_rank, world_size=world_size)
     trainer.setup()
     trainer.train()