fix: avoid duplicate optimizer step with static capture

[SOlo-travelr]

PhysicsNeMo Pull Request

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

CODEOWNERS review map

Current for commit 416ec4169e19. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @dallasfoster — 2 file(s)

• physicsnemo/active_learning/loop.py
• test/active_learning/test_loop.py

⏳ @laserkelvin — 2 file(s)

• physicsnemo/active_learning/loop.py
• test/active_learning/test_loop.py

Comment /codeowners-info to refresh.

[greptile-apps]

Greptile Summary

This PR prevents the active-learning loop from issuing a second optimizer update when static capture already owns backward propagation and optimizer stepping.

• Moves the manual optimizer step into the non-static-capture branch.
• Adds a regression test asserting one deterministic parameter update and one optimizer call on both paths.

Important Files Changed

Filename	Overview
physicsnemo/active_learning/loop.py	Correctly removes the duplicate static-capture update, but allows the scheduler to advance when GradScaler suppresses an optimizer update.
test/active_learning/test_loop.py	Adds focused coverage for exactly one optimizer step and the expected parameter delta on static and ordinary paths.

_{Reviews (1): Last reviewed commit: "fix: avoid duplicate optimizer step with..." | Re-trigger Greptile}

[greptile-apps]

Scheduler advances past skipped updates

When CUDA float16 static capture encounters non-finite gradients, GradScaler.step suppresses the optimizer update but the outer loop still advances lr_scheduler, causing subsequent successful updates to use a learning-rate schedule that is ahead of the actual parameter updates.

Knowledge Base Used: Active Learning Module