🐛[BUG]: DefaultTrainingLoop performs two optimizer steps per minibatch with static capture

Version

main branch

On which installation method(s) does this occur?

Source

Describe the issue

Summary

DefaultTrainingLoop(enable_static_capture=True) performs two optimizer updates for every minibatch.

The static-capture wrapper owns one update through GradScaler.step(optimizer), then DefaultTrainingLoop unconditionally calls optimizer.step() again after the wrapper returns. Because the gradients remain attached, the second call advances the parameters and optimizer state again.

The loop still reports one minibatch and advances the scheduler only once.

Verified on main at commit 043f63bc70dd3d9a491d611b5e27d41af10647b6.

Current code path

1. The active-learning loop wraps the training step in StaticCaptureTraining(model=model, optim=optimizer, ...) when static capture is enabled:

   https://github.com/NVIDIA/physicsnemo/blob/043f63bc70dd3d9a491d611b5e27d41af10647b6/physicsnemo/active_learning/loop.py#L368-L379

2. The wrapper backpropagates and calls self.scaler.step(self.optim):

   https://github.com/NVIDIA/physicsnemo/blob/043f63bc70dd3d9a491d611b5e27d41af10647b6/physicsnemo/utils/capture.py#L157-L174

3. After the wrapped call returns, DefaultTrainingLoop skips only loss.backward() but still calls optimizer.step() unconditionally:

   https://github.com/NVIDIA/physicsnemo/blob/043f63bc70dd3d9a491d611b5e27d41af10647b6/physicsnemo/active_learning/loop.py#L509-L521

Minimal reproduction

This reproduces on CPU; a GPU or active CUDA graph is not required to expose the duplicate optimizer-step ownership.

import torch
from torch.utils.data import DataLoader, TensorDataset

from physicsnemo.active_learning import DefaultTrainingLoop
from physicsnemo.models.mlp import FullyConnected


def run(capture: bool):
    torch.manual_seed(0)
    model = FullyConnected(
        in_features=1,
        out_features=1,
        layer_size=1,
        num_layers=0,
    )

    with torch.no_grad():
        for parameter in model.parameters():
            parameter.fill_(1.0)

    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    loader = DataLoader(
        TensorDataset(torch.ones(1, 1), torch.zeros(1, 1)),
        batch_size=1,
    )

    def step(module, batch):
        x, y = batch
        return (module(x) - y).square().mean()

    before = [parameter.detach().clone() for parameter in model.parameters()]

    DefaultTrainingLoop(
        enable_static_capture=capture,
        use_progress_bars=False,
    )(
        model,
        optimizer,
        loader,
        max_epochs=1,
        train_step_fn=step,
        device=torch.device("cpu"),
        dtype=torch.float32,
    )

    after = [parameter.detach().clone() for parameter in model.parameters()]
    return [(new - old).item() for old, new in zip(before, after)]


print("eager:", run(False))
print("capture:", run(True))

Observed:

eager:   [-0.4, -0.4]
capture: [-0.8, -0.8]

Expected behavior

Static capture should preserve the eager training semantics: exactly one optimizer update per minibatch.

For this deterministic example, the eager and captured parameter deltas should match.

Impact

Every use of DefaultTrainingLoop(enable_static_capture=True) silently applies an extra optimizer update per minibatch. This:

• changes the effective optimization trajectory;
• advances stateful optimizers twice;
• makes scheduler steps, minibatch counts, logs, and training budgets disagree with the actual update count; and
• can produce apparently valid training results using unintended optimization semantics.

Suggested fix

Give optimizer-step ownership to exactly one layer. Either:

• suppress the outer optimizer.step() when StaticCaptureTraining owns backward and step; or
• restructure the capture wrapper so the outer training loop remains the sole optimizer-step owner.

A regression test should run one deterministic minibatch with static capture disabled and enabled, then assert:

1. identical parameter deltas; and
2. exactly one optimizer-step call in both paths.

Minimum reproducible example

Relevant log output

Environment details

Start in physicsnemo/active_learning/loop.py around the static-capture path and inspect physicsnemo/utils/capture.py around StaticCaptureTraining and scaler.step(self.optim). Reproduce the issue with the provided deterministic CPU example, then add a regression test covering eager and captured training. Done means both paths perform exactly one optimizer step and produce identical parameter deltas.