fix: Per-point LR adjustment is not persisted in optimizer state

[andrewwhitecdw]

Problem

fix: Per-point LR adjustment is not persisted in optimizer state

Fix

Replace:

                    scaled_step_size = step_size * per_point_lr
                    p.data.add_(-scaled_step_size * (exp_avg / denom))
                    per_point_lr = self._adjust_per_point_lr(per_point_lr, grad, mask)

with:

                    scaled_step_size = step_size * per_point_lr
                    p.data.add_(-scaled_step_size * (exp_avg / denom))
                    per_point_lr = self._adjust_per_point_lr(per_point_lr, grad, mask)
                    group['per_point_lr'] = per_point_lr

Files changed

• scene/per_point_adam.py