Repository: nerfstudio-project/nerfstudio
Pull request: fix pytorch 2.0 optimizer no inf check error
fix #1883 #1613.

The reason why there's no inf check is because proposal network did not compute gradient at some iteration step. And there is no need to do optimizer.step for proposal network's parameter.

Problem statement:
Pytorch 2.0 nerfacto optimizer inf check error
**Describe the bug**
Nerfacto fails under PyTorch 2.0 with the `no inf checks error`.

**Error**
```
Traceback (most recent call last):
  File "/home/kulhajon/.conda/envs/nerfstudio-torch20/bin/ns-train", line 8, in <module>
    sys.exit(entrypoint())
  File "/home/kulhajon/repos/nerfstudio/scripts/train.py", line 247, in entrypoint
    main(
  File "/home/kulhajon/repos/nerfstudio/scripts/train.py", line 233, in main
    launch(
  File "/home/kulhajon/repos/nerfstudio/scripts/train.py", line 172, in launch
    main_func(local_rank=0, world_size=world_size, config=config)
  File "/home/kulhajon/repos/nerfstudio/scripts/train.py", line 87, in train_loop
    trainer.train()
  File "/home/kulhajon/repos/nerfstudio/nerfstudio/engine/trainer.py", line 236, in train
    loss, loss_dict, metrics_dict = self.train_iteration(step)
  File "/home/kulhajon/repos/nerfstudio/nerfstudio/utils/profiler.py", line 93, in inner
    out = func(*args, **kwargs)
  File "/home/kulhajon/repos/nerfstudio/nerfstudio/engine/trainer.py", line 425, in train_iteration
    self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
  File "/home/kulhajon/repos/nerfstudio/nerfstudio/engine/optimizers.py", line 130, in optimizer_scaler_step_all
    grad_scaler.step(optimizer)
  File "/home/kulhajon/.conda/envs/nerfstudio-torch20/lib/python3.10/site-packages/torch/cuda/amp/grad_scaler.py", line 368, in step
    assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."
AssertionError: No inf checks were recorded for this optimizer.
```

Hints from issue discussion:

Commit message: Merge pull request #2 from ksnzh/ksnzh-patch-1

Patch to review:
diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py
index 3f01fb2e96..a408bcc3a0 100644
--- a/nerfstudio/engine/optimizers.py
+++ b/nerfstudio/engine/optimizers.py
@@ -127,7 +127,8 @@ def optimizer_scaler_step_all(self, grad_scaler: GradScaler) -> None:
             if max_norm is not None:
                 grad_scaler.unscale_(optimizer)
                 torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm)
-            grad_scaler.step(optimizer)
+            if optimizer.param_groups[0]["params"][0].grad is not None:
+                grad_scaler.step(optimizer)
 
     def optimizer_step_all(self) -> None:
         """Run step for all optimizers."""


Code review comments:
diff_hunk: @@ -127,7 +127,8 @@ def optimizer_scaler_step_all(self, grad_scaler: GradScaler) -> None:
             if max_norm is not None:
                 grad_scaler.unscale_(optimizer)
                 torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm)
-            grad_scaler.step(optimizer)
+            if optimizer.param_groups[0]["params"][0].grad is not None:
path: nerfstudio/engine/optimizers.py
text: @user1:
I generally like it. I would just suggest a safer condition:
```
if any(any(p.grad is not None for p in g["params"]) for g in optimizer.param_groups):
```